[AtCoder] ABC 029 D – 1

2020年11月22日2020年12月12日

問題

方針

桁 DP を用いて考えます。\( d(i, j, k) \) を \( i \) 桁目まで調べた時、\( 1 \) の数が \( j \) であり、未満フラグが \( k \) であるものの数とします。\( k = 0 \) であるとき、\( N \) の \( i \) 桁目まで一致している状態で、\( k = 1 \) は \( N \) より小さいことが確定している状態を表します。

コード

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

int main() {
    string N;
    cin >> N;
    int n = N.length();
    ll d[n][n + 1][2]{}; // d[i][j][k] k: 未満フラグ, i桁目までに1がj個である数
    if (N[0] == '1') {
        d[0][1][0] = 1;
        d[0][0][1] = 1;
    } else {
        d[0][0][0] = 1;
        d[0][1][1] = 1;
        d[0][0][1] = N[0] - '0' - 1;
    }
    for (int i = 0; i < n - 1; i++) {
        int t = N[i + 1] - '0';
        // d(i, j, k) を遷移させる
        for (int j = 0; j < n; j++) {
            // d(i, j, 0) の遷移
            if (t == 1) {
                d[i + 1][j + 1][0] += d[i][j][0];
                d[i + 1][j][1] += d[i][j][0]; // i+1桁目: 0
            } else {
                d[i + 1][j][0] += d[i][j][0];
                if (t > 1) {
                    d[i + 1][j][1] += (t - 1) * d[i][j][0]; // i+1桁目: t-1以下の1除く数字にする
                    d[i + 1][j + 1][1] += d[i][j][0]; // i+1桁目: 1
                }
            }
            // d(i, j, 1) の遷移
            d[i + 1][j + 1][1] += d[i][j][1]; // i+1桁目: 1
            d[i + 1][j][1] += d[i][j][1] * 9; // i+1桁目: 1以外
        }
    }
    ll ans = 0;
    for (ll i = 1; i <= n; i++) {
        ans += i * (d[n - 1][i][0] + d[n - 1][i][1]);
    }
    cout << ans << "\n";
    return 0;
}

参考

桁DP(Digit DP) を考え方から問題例まで徹底解説！