読者です 読者をやめる 読者になる 読者になる

pekempeyのブログ

競技プログラミングに関する話題を書いていきます。

桁DP入門

アルゴリズム解説 動的計画法

この記事では桁DPについて以下の問題を取り上げ説明する

A以下の非負整数のうち、「3の倍数または3の付く」かつ「8の倍数でない」数がいくつあるか求めよ

いきなりこの問題を解くのは難しいと思うので、

  1. A以下の非負整数の総数を求める
  2. A以下の非負整数のうち、3の付く数の総数を求める
  3. A以下の非負整数のうち、「3が付くまたは3の倍数」の数の総数を求める
  4. A以下の非負整数のうち、「3が付くまたは3の倍数」かつ「8の倍数でない」数の総数を求める

という4つの問題を順に説明していく。

A以下の非負整数の総数を求める

左から順に数字を決めていこう。
f:id:pekempey:20151208203458p:plain
?の部分に来る可能性のある数字は何だろうか。 ?は0~9のどれでもいい。なぜなら上から2桁目がAより小さいため、?に何を選んでもA以上にはならないからである。

次の場合はどうだろうか。
f:id:pekempey:20151208203812p:plain
この場合?に持ってこれるのは0~6の数字のみである。

この2つの違いは、今見てる桁よりも前の桁に「Aより小さいものがあるかどうか」だけ。だから、具体的に何を決めたかが分からなくても「A未満であることが確定している」ことだけ分かればいい。
f:id:pekempey:20151208204204p:plain
less=1なら?に来るのは0~9の数字。

f:id:pekempey:20151208204130p:plain
less=0なら?に来るのは0~6の数字。?に0~5の数字を選ぶとless=1になる。

dp[i][j]
i: 上からi桁目まで見ている
j: A未満であることが確定している
#include <iostream>
#include <string>
#define rep(i, a) for (int i = 0; i < (a); i++)
using namespace std;

const int mod = 1e9 + 7;
int dp[101010][2]; // pos, less

int main() {
    string A;
    cin >> A;
    int n = A.length();

    dp[0][0] = 1;
    rep (i, n) rep (j, 2) {
        int lim = j ? 9 : A[i] - '0';
        rep (d, lim + 1) {
            (dp[i + 1][j || d < lim] += dp[i][j]) %= mod;
        }
    }

    int ans = 0;
    rep (j, 2) (ans += dp[n][j]) %= mod;
    cout << ans << endl;
    return 0;
}

A以下の非負整数のうち、3の付く数の総数を求める

「A未満であることが確定している」という状態だけでなく「すでに3を選んでいる」という状態も付け加えればいい。?で3を選ぶとhas3=1になる。

f:id:pekempey:20151208205009p:plain

dp[i][j][k]
i: 上からi桁目まで見ている
j: A未満であることが確定している
k: すでに3を選んでいる
#include <iostream>
#include <string>
#define rep(i, a) for (int i = 0; i < (a); i++)
using namespace std;

const int mod = 1e9 + 7;
int dp[101010][2][2]; // pos, less, has3

int main() {
    string A;
    cin >> A;
    int n = A.length();

    dp[0][0][0] = 1;
    rep (i, n) rep (j, 2) rep (k, 2) {
        int lim = j ? 9 : A[i] - '0';
        rep (d, lim + 1) {
            (dp[i + 1][j || d < lim][k || d == 3] += dp[i][j][k]) %= mod;
        }
    }

    int ans = 0;
    rep (j, 2) ans += dp[n][j][1];
    cout << ans << endl;
    return 0;
}

A以下の非負整数のうち、「3が付くまたは3の倍数」の数の総数を求める

「A未満であることが確定している」、「すでに3を選んでいる」という状態だけでなく、「mod 3の値」も状態に組み込めばいい。

f:id:pekempey:20151208210016p:plain

dp[i][j][k][l]
i: 上からi桁目まで見ている
j: A未満であることが確定している
k: すでに3を選んでいる
l: mod 3の値

「3が付くまたは3の倍数」という条件はDPの更新段階では現れず、ans計算時に現れる。 次のmod 3の値は(l+?) mod 3で計算できる。

#include <iostream>
#include <string>
#define rep(i, a) for (int i = 0; i < (a); i++)
using namespace std;

const int mod = 1e9 + 7;
int dp[101010][2][2][3]; // pos, less, has3, mod3

int main() {
    string A;
    cin >> A;
    int n = A.length();

    dp[0][0][0][0] = 1;
    rep (i, n) rep (j, 2) rep (k, 2) rep (l, 3) {
        int lim = j ? 9 : A[i] - '0';
        rep (d, lim + 1) {
            (dp[i + 1][j || d < lim][k || d == 3][(l + d) % 3] += dp[i][j][k][l]) %= mod;
        }
    }

    int ans = 0;
    rep (j, 2) rep (k, 2) rep (l, 3) {
        if (k || l == 0)  {
            (ans += dp[n][j][k][l]) %= mod;
        }
    }
    cout << ans << endl;
    return 0;
}

A以下の非負整数のうち、「3が付くまたは3の倍数」かつ「8の倍数でない」数の総数を求める

「mod 8の値」も状態に組み込む。

f:id:pekempey:20151208210028p:plain

dp[i][j][k][l][m]
i: 上からi桁目まで見ている
j: A未満であることが確定している
k: すでに3を選んでいる
l: mod 3の値
m: mod 8の値

次のmod 8の値は(10*m+?) mod 8で計算できる。

#include <iostream>
#include <string>
#define rep(i, a) for (int i = 0; i < (a); i++)
using namespace std;

const int mod = 1e9 + 7;
int dp[101010][2][2][3][8]; // pos, less, has3, mod3, mod8

int main() {
    string A;
    cin >> A;
    int n = A.length();

    dp[0][0][0][0][0] = 1;
    rep (i, n) rep (j, 2) rep (k, 2) rep (l, 3) rep (m, 8) {
        int lim = j ? 9 : A[i] - '0';
        rep (d, lim + 1) {
            (dp[i + 1][j || d < lim][k || d == 3][(l + d) % 3][(m * 10 + d) % 8] += dp[i][j][k][l][m]) %= mod;
        }
    }

    int ans = 0;
    rep (j, 2) rep (k, 2) rep (l, 3) rep (m, 8) {
        if ((k || l == 0) && m != 0)  {
            (ans += dp[n][j][k][l][m]) %= mod;
        }
    }
    cout << ans << endl;
    return 0;
}

一番最後に説明した問題は、yukicoderの世界のなんとか3とほとんど同じもの。

ある芸人は「3の倍数 もしくは 3のつく数」の時「アホ」になり、「8の倍数」の時「青春」するという。
A 以上 B 以下の整数のうち、「アホ」になりかつ「青春」しない数がいくつあるかを 109+7 で割った時の余りで求めてください。

この問題の場合、A以上B以下のものを求める必要があるが、

f(n):=n以下の条件を満たす非負整数の総数

とすればf(B)-f(A-1)で求められるので、あまり気にする必要はない。

A-1を計算してもいいが、lessフラグが0のものを無視すればA未満の総数が得られる。

例題

次の2つも桁DPだけど今回説明した桁DPとは違って、桁上げや桁借りを使って下から決めていくタイプ。

ひとこと

僕がrepマクロ使い始めたのは桁DPのせい。それまではスニペット使ってたんだけど、6重のforは見た目が怖いからやめた。C++はDPの添字に論理演算の結果をそのまま使えるので少し楽。