【TDPC】E - 数

atcoder.jp

桁DPの問題。初めて桁DPについて学んだので記事にします。

桁DPとは？

・N以下の自然数で条件を満たすものは何通りか？
・N以下の自然数で条件を満たすものの中で最大のものはいくつか？
みたいな問題で活躍するアルゴリズムです。

説明は少し複雑になるので以下の記事たちに丸投げします。わかりやすいです。

qiita.com

drken1215.hatenablog.com

アルゴリズムロジック- 桁DP(Digit DP) を考え方から問題例まで徹底解説！

解法

DPテーブルの中身は、条件を満たす場合の数にしたいので、添字に保持する情報を考えます。
条件は、Dの倍数であるかどうかなので、mod Dで考えて問題ありません。
Dは100以下の値しか取らないので、Nの桁数×Dでもせいぜい10⁶にしかならないため、mod Dで考えた各桁の和の情報を添字に保持することにします。
よって以下のDPを考えることができます。

dp[i][smaller][j] ← i: 上から見た桁数, smaller: 未満フラグ, j: 各桁の和(mod D)

初期条件は、
dp[0][0][0] ← 1
とします。実際には存在しないですが、0桁目に0があると考えると上の初期条件になります。
実際には存在しないので、最後に答えから1を引く必要があります。

桁DPの遷移方法には決まりがあるので、それについては上の記事をご覧ください。
今回は次のように遷移します。
ここで、i桁目の数字をn[i]とします。

• 未満フラグが立っていない場合
  → 0からn[i]までがN以下になるので遷移可能、n[i]未満の場合は未満フラグをつける
• 未満フラグが立っている場合
  → 0から9のどの値をとっても必ずN以下になるので遷移可能、もちろん未満フラグはつけっぱなし

これをD未満の全ての自然数に対して行い、遷移完了です。

そして、答えは
dp[N.size()][0][0] + dp[N.size()][1][0] - 1
になります。必ず1を引くようにしましょう。

実装

遷移は、論理記号をうまく使うことで1行で表すことができます。

void solve(long long D, std::string N) {
    int l = N.size();

    // 文字列から数字の配列に変換
    vector<int> n(l);
    REP (i, l) n[i] = N[i] - '0';

    // DPテーブル
    vector<vector<vector<mint>>> dp(l+1, vector<vector<mint>>(2, vector<mint>(D, 0)));

    // 初期値
    dp[0][0][0] = 1;

    // 桁DP
    REP (i, l) {
        REP (smaller, 2) {
            for (int x = 0; x <= (smaller ? 9 : n[i]); x++) {
                REP (j, D) {
                    dp[i + 1][smaller || x < n[i]][(j + x) % D] += dp[i][smaller][j];
                }
            }
        }
    }

    c(dp[l][0][0] + dp[l][1][0] - 1)
}