BEN2のブログ

たまに書いています。

ABC 154 E - Almost Everywhere Zero

競技プログラミング

ABC 154 E - Almost Everywhere Zero に関連する、桁 DP についてのメモを残します。

参考

桁 DP の前に

314 を超えない非負整数の個数を求める問題を考えてみる。
百の位、十の位、一の位の順に数字を決めていくことにする。001 は 1 とみなす。

まず、百の位に入れられる数字は 0, 1, 2, 3 の 4 つに限られる。
ここで注目しておきたいのは、百の位に 3 を入れるかどうかで、十の位に入れることのできる数字が変わってくる点。

百の位に 0, 1, 2 のいずれかを入れた場合、この時点で 314 を超えないことが確定する (この状態を「未満フラグが立っている」と呼ぶことにする) ため、このあとの十の位に入れる数字は自由に決めることができる。
さらに一の位に入れる数字も自由に決められることからわかるように、未満フラグは一度立てば、その後フラグが折れることはない。

一方、百の位に 3 を入れた場合、この時点では「確実に 314 を超えない」とは言えない (未満フラグは立たない)。十の位に入れられる数字は 0, 1 の 2 つに限られる。
このあと、もし十の位に 0 を入れたら、未満フラグが立つ。

ちなみに、百の位に入れる数字が 4 つに限られていたのは、もちろん未満フラグが立っていなかったためである。

遷移図は次の通り。青は未満フラグが立っておらず、橙は未満フラグが立っている。

f:id:BEN2suzuka:20200212025442p:plain

桁 DP

上の遷移図を下図のようにまとめ、

{\rm{dp}} [ i ] [ j ] := 左から i 桁決めて、 j = 1 ならば未満フラグが立っている

とする。

f:id:BEN2suzuka:20200212163520p:plain f:id:BEN2suzuka:20200212025504p:plain

先ほどと同じように、314 を超えないように百の位から順に数字を決めていく。未満フラグは立っていないため、百の位は 0, 1, 2, 3 の 4 つに限られる。0, 1, 2 を入れる場合、未満フラグが立つので、 {\rm{dp}} [ 0 ] [ 0 ] を {\rm{dp}} [ 1 ] [ 1 ] に配る。3 を入れる場合、未満フラグは立たないので、 {\rm{dp}} [ 1 ] [ 0 ] に配る。このときの擬似コードを下に示す。

f:id:BEN2suzuka:20200212164146p:plain

Str = "314"
for d = 0 to Str[0]  // d は 1 桁目に入れる数字
  if d < Str[0]  // 未満フラグが立つ
    dp[1][1] += dp[0][0]
  else
    dp[1][0] += dp[0][0]

百の位に 0, 1, 2 のいずれかを入れた場合、未満フラグが立っているので、このあとの十の位は自由に決められる。このときの擬似コードは以下の通り。 {\rm{dp}} [ 1 ] [ 1 ] を {\rm{dp}} [ 2 ] [ 1 ] に 10 回配っている。

f:id:BEN2suzuka:20200212163614p:plain

for d = 0 to 9  // 十の位は自由
  dp[2][1] += dp[1][1]

非負整数 N を超えない非負整数の個数を求めるコードは次の通り。

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

int main() {
  int N; cin >> N;
  string S = to_string(N);
  int L = S.size();
  // dp[i][j] := 左から i 桁決めて、j = 1 のとき未満フラグが立っている
  vector<vector<int>> dp(L+1, vector<int>(2));
  dp.at(0).at(0) = 1;  // dp[0][1] は 0
  // 左から i+1 桁目の数字を決める
  for (int i = 0; i < L; i++) {
    // 未満フラグが立っていない場合 (dp[i][0] からの遷移)
    int D = S.at(i) - '0';  // 整数 N の左から i+1 桁目の数字
    for (int d = 0; d <= D; d++) {
      if (d < D) {  // 未満フラグが立つ
        dp.at(i+1).at(1) += dp.at(i).at(0);
      }
      else {
        dp.at(i+1).at(0) += dp.at(i).at(0);
      }
    }
    // 未満フラグが立っている場合 (dp[i][1] からの遷移)
    for (int d = 0; d <= 9; d++) {
      dp.at(i+1).at(1) += dp.at(i).at(1);
    }
  }
  cout << dp.at(L).at(0) + dp.at(L).at(1) << endl;
}

洗練されたコードは、

この記事がとても参考になります。

{\rm{dp}} を 2 つに分ける

  • {\rm{dp0}} [ i ] := 左から i 桁決めて、未満フラグは立っていない
  • {\rm{dp1}} [ i ] := 左から i 桁決めて、未満フラグは立っている

とすることもできます。

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

int main() {
  int N; cin >> N;
  string S = to_string(N);
  int L = S.size();
  vector<int> dp0(L+1), dp1(L+1);
  dp0.at(0) = 1;
  for (int i = 0; i < L; i++) {
    // 未満フラグが立っていない場合 (dp0[i] からの遷移)
    int D = S.at(i) - '0';
    for (int d = 0; d <= D; d++) {
      if (d < D) {  // 未満フラグが立つ
        dp1.at(i+1) += dp0.at(i);
      }
      else {
        dp0.at(i+1) += dp0.at(i);
      }
    }
    // 未満フラグが立っている場合 (dp1[i] からの遷移)
    for (int d = 0; d <= 9; d++) {
      dp1.at(i+1) += dp1.at(i);
    }
  }
  cout << dp0.at(L) + dp1.at(L) << endl;
}

このあと、この方法で ABC 154 E - Almost Everywhere Zero を解いてみる。

0 でない数字を K 個含む整数の個数

0 でない数字を何個含んでいるのか、という情報が必要なので、

  • {\rm{dp}} [ i ] [ j ] [ k ] := 左から i 桁決めて、 j = 1 ならば未満フラグが立っている。0 でない数字を k 個含む

が考えられる。今回は 2 つに分けて、

  • {\rm{dp0}} [ i ] [ j ] := 左から i 桁決めて、0 でない数字を j 個含む。未満フラグは立っていない
  • {\rm{dp1}} [ i ] [ j ] := 左から i 桁決めて、0 でない数字を j 個含む。未満フラグは立っている

としてみる。
左から i+1 桁目の数字を決めるとき、未満フラグが立っていない場合 ( {\rm{dp0}} [ i ] [ j ] からの遷移) は、

Str = "31415"
for j = 0 to 3  // 0 でない数字を何個含むか
  for d = 0 to Str[i]  // d は i+1 桁目に入れる数字
    nj = j
    if d != 0
      nj++  // 0 でない数字が 1 個増える
    if d < Str[i]  // 未満フラグが立つ
      dp1[i+1][nj] += dp0[i][j]
    else
      dp0[i+1][nj] += dp0[i][j]

このような擬似コードで、未満フラグが立っている場合 ( {\rm{dp1}} [ i ] [ j ] からの遷移) は、

for j = 0 to 3
  for d = 0 to 9
    nj = j
    if d != 0
      nj++
    dp1[i+1][nj] += dp1[i][j]

となる。

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

int main() {
  string S; cin >> S;
  int K; cin >> K;
  int L = S.size();
  // dp[i][j] := 左から i 桁決めて、0 でない数字を j 個含む
  vector<vector<long long>> dp0(L+1, vector<long long>(5));
  vector<vector<long long>> dp1(L+1, vector<long long>(5));
  dp0.at(0).at(0) = 1;
  for (int i = 0; i < L; i++) {
    // j = 4, 5, ... からの遷移は不要
    for (int j = 0; j < 4; j++) {
      // 未満フラグが立っていない場合 (dp0[i][j] からの遷移)
      int D = S.at(i) - '0';
      for (int d = 0; d <= D; d++) {
        int nj = j;
        if (d != 0) nj++;
        if (d < D) {  // 未満フラグが立つ
          dp1.at(i+1).at(nj) += dp0.at(i).at(j);
        }
        else {
          dp0.at(i+1).at(nj) += dp0.at(i).at(j);
        }
      }
      // 未満フラグが立っている場合 (dp1[i][j] からの遷移)
      for (int d = 0; d <= 9; d++) {
        int nj = j;
        if (d != 0) nj++;
        dp1.at(i+1).at(nj) += dp1.at(i).at(j);
      }
    }
  }
  cout << dp0.at(L).at(K) + dp1.at(L).at(K) << endl;
}

最後に、今回の問題で指定された探索範囲が 0 以上ではなく 1 以上であることに注意する。
このコードでは、整数 0 も探索範囲に含まれており、整数 0 は dp1[L][0] に入る。
もし K = 0 が与えられた場合、0 以外の数字を 1 個も使わない 1 以上 N 以下の整数は存在しないので、0 を出力しなければならないが、このコードでは 1 と出力してしまう。
今回は K = 0 が与えられることはないため、このことを意識しなくても AC できる (コンテスト中、気づきませんでした)。