桁 DP の問題では,「解法が分かること」と「バグりにくい実装ができること」の間に大きな乖離があると思います.
この記事では,「解法は分かるけど実装が破滅する…」という方に向けて,バグりにくい実装方針を紹介します.AtCoder 水色〜青色程度の方を想定読者としています.
実装方針
結論ですが,桁 DP は以下の方針に従うとほぼ機械的に書け,バグりにくいと思います:
- 「配る DP」で書く
- 「遷移前の状態」「状態遷移」「遷移後の状態」を変数で管理し,遷移前の状態と状態遷移の列挙を全て
for
文に任せる for
文の中では,まず遷移後の状態を表す変数を遷移前の状態で初期化し,必要に応じて遷移後の状態を書き換える- 不要な遷移は
continue
で抜ける
この実装方針だと,+=
や chmax
を含む「遷移式」を 1 本しか書かずに済み,コードが簡潔になります.
以下では例題とともに実装を紹介します.全て 0-indexed に統一して書きます.
例題1: ABC154: E - Almost Everywhere Zero
問題ページ: ABC154: E - Almost Everywhere Zero
$N$ の桁数を $n$ とします.$N$ 以下の非負整数を(必要なら $0$ 埋めすることで)$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
問題ページ: 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 \geq L,$ $y \leq 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]; // 遷移式
}
}
→ 提出コード
実装のポイント
- 遷移後の状態を表す変数を遷移前の状態で初期化し(5 行目),状態の更新が必要な場合のみ書き換えます(10,14,18 行目)
- 不要な遷移は continue で抜けます(7,9,13,17行目)
- 遷移式を 1 本しか書いていません(20 行目)
- 12 〜 18 行目の処理は例題 1 と全く同様です
余談
以前 Coincidence の実装で地獄を味わったので,この記事を書きました.桁 DP の実装で破滅する人が一人でも減れば幸いです.
参考ツイート:
ABC138 F - Coincidence を解くのに地獄を体現したかのようなソースコードを生成してしまった (無事ACできて良かった...) pic.twitter.com/1ufajyikW2
— opt (@opt_cp) June 28, 2020