桁 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]; // 遷移式
  }
}

提出コード

実装のポイント

  • 遷移前の状態を変数 jk で,状態遷移を変数 x で, 遷移後の状態を変数 j2k2 で表しています
  • 遷移後の状態を表す変数を遷移前の状態で初期化し(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]; // 遷移式
  }
}

提出コード

実装のポイント

  • 遷移後の状態を表す変数を遷移前の状態で初期化し(5 行目),状態の更新が必要な場合のみ書き換えます(10,14,18 行目)
  • 不要な遷移は continue で抜けます(7,9,13,17行目)
  • 遷移式を 1 本しか書いていません(20 行目)
  • 12 〜 18 行目の処理は例題 1 と全く同様です

余談

以前 Coincidence の実装で地獄を味わったので,この記事を書きました.桁 DP の実装で破滅する人が一人でも減れば幸いです.

参考ツイート:

タイトルとURLをコピーしました