【ABC154】E - Almost Everywhere Zero

atcoder.jp

解法

参考: 桁DPについて
nizi-24.hatenablog.com

桁DPをしよう。
dp[i][smaller][k] ← i: 上から何桁目まで確定しているか、smaller: 未満フラグ、k: 今までに何個0以外の数字が登場したか

遷移するときには、0かどうかで場合分けすればよい。
ここで、Nの上からi桁目はn[i]とする。

• i桁目にx = 0を採用する場合
  • dp[i][smaller || x < n[i]][j] += dp[i][smaller][j]
• i桁目にx ≠ 0を採用する場合
  • dp[i][smaller || x < n[i]][j + 1] += dp[i][smaller][j]

という感じで遷移すれば良い。
答えは、dp[N.size()][0][K] + dp[N.size()][1][K] になる。

実装

long long K;
string N;
void solve() {
    int l = N.size();
    // DPテーブル
    vector<vector<vector<long long>>> dp(l+1, vector<vector<long long>>(2, vector<long long>(K+2, 0)));

    // 初期条件
    dp[0][0][0] = 1;

    // 遷移
    REP (i, l) {
        REP (smaller, 2) {
            for (int x = 0; x <= (smaller ? 9 : N[i] - '0'); x++) {
                REP (j, K+1) {
                    if (x == 0) dp[i + 1][smaller || x < (N[i] - '0')][j] += dp[i][smaller][j];
                    else dp[i + 1][smaller || x < (N[i] - '0')][j + 1] += dp[i][smaller][j];
                }
            }
        }
    }

    cout << dp[l][0][K] + dp[l][1][K] << endl;
}