AtCoder Educational DP Contest S問題解いてみた

今回の問題

問題文

$1$ 以上 $K$ 以下の整数のうち、十進表記における各桁の数字の総和が $D$ の倍数であるようなものは何個でしょうか？
$10^9 + 7$ で割った余りを求めてください。

制約

• 入力はすべて整数である。
• $1 \leq K < 10^{10000}$
• $1 \leq D \leq 100$

回答

#include <bits/stdc++.h>
using namespace std;

const int MOD = 1e9 + 7;

// dp[桁数][未満フラグ][余り]
// 最大10000桁なので、メモリ制限に注意しつつ確保
long long dp[10005][2][100];

int main() {
    string K;
    int D;
    cin >> K >> D;

    int n = K.size();

    // 初期化
    // 0桁目、まだKと一致(smaller=0)、余り0 の状態が1通り
    dp[0][0][0] = 1;

    for (int i = 0; i < n; i++) {
        // 文字列Kのi文字目を数値に変換
        int limit = K[i] - '0';

        for (int smaller = 0; smaller < 2; smaller++) {
            for (int rem = 0; rem < D; rem++) {
                // その状態に到達不可能な場合はスキップ
                if (dp[i][smaller][rem] == 0) continue;

                // 次の桁に置ける数字 x の範囲を決める
                // smaller=1 (すでにK未満確定) なら 0~9
                // smaller=0 (Kと一致中) なら 0~limit
                int upper = (smaller == 1) ? 9 : limit;

                for (int x = 0; x <= upper; x++) {
                    // 次の smaller 状態
                    // すでに smaller=1 なら 1 のまま
                    // まだ smaller=0 でも、今回 limit より小さい数(x < limit)を選べば 1 になる
                    int next_smaller = smaller;
                    if (smaller == 0 && x < limit) {
                        next_smaller = 1;
                    }

                    // 次の余り状態
                    int next_rem = (rem + x) % D;

                    // 遷移
                    dp[i + 1][next_smaller][next_rem] = (dp[i + 1][next_smaller][next_rem] + dp[i][smaller][rem]) % MOD;
                }
            }
        }
    }

    // 答えは N桁目まで見たときの、smaller=0と1両方の、余りが0のもの
    long long ans = (dp[n][0][0] + dp[n][1][0]) % MOD;

    // 0を含んでしまっているので1引く
    // (ただし、負の数にならないようにMODを足してから割る)
    ans = (ans - 1 + MOD) % MOD;

    cout << ans << endl;

    return 0;
}

自分なりの解説

dp[i][smaller][rem]を上からi桁目、smallerで上限の制約を受けるかどうか、remで余りを設定している。上からi桁目まではわかったが、smallerとremは回答を見ないとわからなかった。