桁 DP の簡潔な実装

桁 DP の問題では，「解法が分かること」と「バグりにくい実装ができること」の間に大きな乖離があると思います．

この記事では，「解法は分かるけど実装が破滅する…」という方に向けて，バグりにくい実装方針を紹介します．AtCoder 水色〜青色程度の方を想定読者としています．

実装方針

結論ですが，桁 DP は以下の方針に従うとほぼ機械的に書け，バグりにくいと思います:

「配る DP」で書く
「遷移前の状態」「状態遷移」「遷移後の状態」を変数で管理し，遷移前の状態と状態遷移の列挙を全て for 文に任せる
for 文の中では，まず遷移後の状態を表す変数を遷移前の状態で初期化し，必要に応じて遷移後の状態を書き換える
不要な遷移は continue で抜ける

この実装方針だと，+= や chmax を含む「遷移式」を 1 本しか書かずに済み，コードが簡潔になります．

以下では例題とともに実装を紹介します．全て 0-indexed に統一して書きます．

例題1: ABC154: E - Almost Everywhere Zero → 問題

$N$ の桁数を $n$ とします．$\mathtt{dp}[i][j][k]$ を，$n$ 桁の非負整数 $X$ の上位 $i$ 桁の決め方であって，条件

非ゼロの桁がちょうど $j$ 個ある
$X$ が $N$ より小さいことが確定 $\{$しない $(k=0), $ する $(k=1)$$\}$

をともに満たすものの数，と定めます．これは $i = 0,1,\dots$ と順に計算できます．

実装は以下のとおりです:

rep(i, n) {
  int Ni = N[i] - '0'; // N の上から i 桁目の数
  rep(j, K+1) rep(k, 2) rep(x, 10) {
    int j2 = j, k2 = k; // 遷移後の状態を表す変数 j2, k2 を，遷移前の状態で初期化

    if (x != 0) ++j2;     // X の上から i 桁目の数字 x が非ゼロなら j2 を更新
    if (j2 > K) continue; // 非ゼロの桁が K より多いなら無視

    if (!k && (x > Ni)) continue; // X が N より大きいことが確定したら無視
    if (x < Ni) k2 = 1;           // X が N より小さいことが確定したら k2 を更新

    dp[i+1][j2][k2] += dp[i][j][k]; // 遷移式
  }
}

→ 提出コード

実装のポイント

遷移前の状態を変数 j，k で，状態遷移を変数 x で，遷移後の状態を変数 j2，k2 で表しています
遷移後の状態を表す変数を遷移前の状態で初期化し（4 行目），状態の更新が必要な場合のみ書き換えます（6，10 行目）
不要な遷移は continue で抜けます（7，9 行目）
遷移式を 1 本しか書いていません（12 行目）
「$0$ 以上 $N$ 以下の整数」という条件は，9 〜 10 行目のように処理できます．頻出なので覚えてしまうとよいでしょう

例題2: ABC138: F - Coincidence → 問題

$x,y$ を $M = 60$ 桁の二進表記したときの上から $i$ 桁目の数字を $x_i, y_i$ と書きます．公式解説にもある通り，問題を言い換えると，整数の組 $(x, y)$ で，条件

任意の $i$ に対し $(x_i, y_i) \neq (1, 0)$
$(x_0, y_0), (x_1, y_1),\dots$ に $(1, 1)$ より先に $(0, 1)$ が現れない
$x \geq L$
$y \leq R$

を全て満たすものを数え上げればよいことが分かります．

そこで $\mathtt{dp}[i][j][l][r]$ を，$x$ と $y$ の上位 $i$ 桁の決め方であって，条件

任意の $k \in [0, i)$ に対し $(x_k, y_k) \neq (1, 0)$
$(x_0, y_0), (x_1, y_1),\dots, (x_{i-1}, y_{i-1})$ に $(1, 1)$ より先に $(0, 1)$ が現れない
$(x_0, y_0), (x_1, y_1),\dots, (x_{i-1}, y_{i-1})$ に $(1, 1)$ が $\{$現れない $(j=0), $ 現れる $(j=1)$$\}$
$x > L$ であることが確定 $\{$しない $(l=0), $ する $(l=1)$$\}$
$y < R$ であることが確定 $\{$しない $(r=0), $ する $(r=1)$$\}$

を全て満たすものの数，と定めます．これは $i = 0,1,\dots$ と順に計算できます．

実装は以下です:

rep(i, M) {
  int Li = L>>(M-1-i)&1; // L の上から i 桁目
  int Ri = R>>(M-1-i)&1; // R の上から i 桁目
  rep(j, 2) rep(l, 2) rep(r, 2) rep(xi, 2) rep(yi, 2) {
    int j2 = j, l2 = l, r2 = r; // 遷移前の状態で初期化

    if (xi > yi) continue; // (xi, yi) = (1, 0) なら無視

    if (!j && (xi < yi)) continue; // (1, 1) が現れる前に (0, 1) が現れたら無視
    if (xi & yi) j2 = 1;           // (1, 1) が現れたら j2 を更新

    // 条件 x ≧ L の処理
    if (!l && (xi < Li)) continue;
    if (xi > Li) l2 = 1;

    // 条件 y ≦ R の処理
    if (!r && (yi > Ri)) continue;
    if (yi < Ri) r2 = 1;

    dp[i+1][j2][l2][r2] += dp[i][j][l][r]; // 遷移式
  }
}

→ 提出コード