Atcoder Begginer Contest 029_D_1

水色の問題.桁DPの立て方をしっかり学べてよい.

問題の説明

$1 \le k \le N$ を満たすすべての整数 $k$ に含まれる $'1'$ の個数はいくつか答えよ.

制約

$1 \le N \lt 10^{9}$

考えた事

まず $N$ の制約に注目したい.上限が $10^{9}$ であり, $k$ を $1$ から $N$ まで実際に動かして $1$ の個数を数え上げる $O(N \log k)$ の解法は厳しく見える.こういった場合に有効なのが桁DPである.

桁DPについては garakutagoya.hatenablog.com

でも記載した通りである. $dp_{i,j}:=(i桁目まで見た時にN未満確定かどうか(jはbool値))$ をベースに条件を付随する.

今回もそれを使う.こうする事で計算量を桁数まで落とせる.計算量を落とせる理由としては $N=352$ である場合に, $10,12$ を $10,12$ として記録するのではなく,「 $2$ 桁目まで見た時に $N$ 未満確定かつ $1$ の個数が $1$ 個」として一緒くたの遷移元として扱えるからである.一般に配列のデータへのアクセスは $O(1)$ で行えるため,要素数を減らせばそれだけ計算量を落とせるという算段である.¹ という事であらためてDP式は $dp_{i,j,k}:=(i桁目まで見た時にN未満確定かどうか(jはbool値)かつ1の個数がk個ある)$ となる.後は遷移の仕方を考えればいい.

桁DPの遷移のさせ方は動かす数と注目している数との大小関係によって場合分けすると見通しがいい.

rep(i, num.size()) {
        ll kazu = num[i] - '0';
        rep(p, 10) {
            if (p < kazu) {
                rep(j, 2) {
                    rep(k, 12) {
                        if (p == 1) {
                            dp[i + 1][1][k + 1] += dp[i][j][k];
                        } else {
                            dp[i + 1][1][k] += dp[i][j][k];
                        }
                    }
                }
            }

上のコードの断片を見てもらいたい.数の取り方は $0~9$ で $10$ 通りあるが,それと $N$ のそれぞれの桁の値とを順に比較している.² この場合は $p$ が $N$ のある桁より小さい時を表している.

この時,この桁より下の桁がどんなものであっても,それが $N$ 以上にはならない.たとえば $N=1287$ の時, $11...$ となる数は下二桁が何であってもこれは $N=1287$ に到達しないのだ.したがって遷移先は $j=1(N未満確定)$ となる.

$k$ については $p=1$ の時は $1$ の個数がひとつ増えるため,それに気をつければいい.

else if (p == kazu) {
                rep(j, 2) {
                    rep(k, 12) {
                        if (p == 1) {
                            dp[i + 1][j][k + 1] += dp[i][j][k];
                        } else {
                            dp[i + 1][j][k] += dp[i][j][k];
                        }
                    }
                }

$p$ と $N$ の注目している桁の数字が同じ時はどうすればいいか.この場合は $p$ に $N未満確定$ させる力はないため,前の数字に依存する.例を挙げると $N=2342$ のとき, $13...$ は $N$ 未満確定だが $23...$ が $N$ 未満確定とはならないという事だ.したがって遷移先も前のものが $N未満確定かどうか$ に依存する.

else {
                rep(k, 12) {
                    if (p == 1) {
                        dp[i + 1][1][k + 1] += dp[i][1][k];
                    } else {
                        dp[i + 1][1][k] += dp[i][1][k];
                    }
                }
            }

$p$ が $N$ の注目している桁の数字より大きい場合はどうすればいいか.この時は遷移元を絞る必要がある. $N=236534$ のとき, $24...$ となるものは有り得ない.したがって, $14...$ となるもののみを考えればいい.このとき明らかに遷移元,遷移先共に $N$ 未満確定となる.

以上のように $dp_{i,j,k}$ を操作する事で解が得られる.聞かれているものは $1$ の個数であるから,答えの計算は $k×(dp_{ketasu,1,k}+dp_{ketasu,0,k})$ を足し合わせるといいだろう.ここまでお読みいただきありがとうございました.

プログラム

#include <bits/stdc++.h>
using namespace std;
#include <atcoder/all>
using namespace atcoder;
using ll = long long;
using vll = vector<long long>;
using vvll = vector<vll>;
using vvvll = vector<vvll>;
#define rep(i, n) for (ll i = (0); i < (ll)(n); i++)
int main() {
    ll N;
    cin >> N;
    string num = to_string(N);
    vvvll dp(15, vvll(2, vll(15, 0)));
    dp[0][0][0] = 1;
    rep(i, num.size()) {
        ll kazu = num[i] - '0';
        rep(p, 10) {
            if (p < kazu) {
                rep(j, 2) {
                    rep(k, 12) {
                        if (p == 1) {
                            dp[i + 1][1][k + 1] += dp[i][j][k];
                        } else {
                            dp[i + 1][1][k] += dp[i][j][k];
                        }
                    }
                }
            } else if (p == kazu) {
                rep(j, 2) {
                    rep(k, 12) {
                        if (p == 1) {
                            dp[i + 1][j][k + 1] += dp[i][j][k];
                        } else {
                            dp[i + 1][j][k] += dp[i][j][k];
                        }
                    }
                }
            } else {
                rep(k, 12) {
                    if (p == 1) {
                        dp[i + 1][1][k + 1] += dp[i][1][k];
                    } else {
                        dp[i + 1][1][k] += dp[i][1][k];
                    }
                }
            }
        }
    }
    ll ans = 0;
    rep(j, 2) {
        rep(k, num.size() + 1) { ans += k * dp[num.size()][j][k]; }
    }
    cout << ans << endl;
}

蟻本の個数制限付き部分和問題の箇所の「bool値を載せるだけのDPにはむだが多い」という話に近い.↩
$numはNをto_string函数で文字列化したもの$ ↩