[AtCoder] ABC 154 E – Almost Everywhere Zero
問題
https://atcoder.jp/contests/abc154/tasks/abc154_e方針
見るからに桁DPの問題なので動的計画法を行います。
\( d(i, j, 0) \) を \( N \) の上位 \( i \) 桁目まで一致する数字で、\( 0 \) でない数字の個数が \( j \) 個であるものが存在するかどうかの値とします。存在するときは \( d(i, j, 0) = 1 \) であり、存在しないときは \( d(i, j, 0) = 0\) となります。次に、\( d(i, j, 1) \) を \( N \) の \( i \) 桁目未満の数字の中で、\( 0 \) でない数字の個数が \( j \) 個であるものの個数とします。
ここで、\( N \) の上位 \( i \) 桁目の数字を \( n_i \) とします。初期値は、
\begin{eqnarray}
d(1, 0, 1) &=& 1\\
d(1, 1, 0) &=& 1\\
d(1, 1, 1) &=& n_1 – 1
\end{eqnarray}
となります。その他の \( d(i, j, k) \) は \( 0 \) で初期化します。次に、\( 2 \) 桁目以降の更新式を考えます。
\( n_i = 0\) のとき
このとき、\( N \) に上位 \( i + 1\) 桁目まで一致する数字は \( i \) 桁目まで一致する状態から動かないので、
\[ d(i + 1, j, 0) \leftarrow d(i, j, 0) \ (0 \leq j \leq 3) \]
となります。\( d(i + 1, j, 1) \) の更新は、\( i + 1 \) 桁目を \( 0 \) にする場合を考え、
\[ d(i + 1, j, 1) \leftarrow d(i, j, 1) \ (0 \leq j \leq 3) \]
とします。また、\( 0 \) 以外の値にすることを考えると、
\[ d(i + 1, j + 1, 1) \leftarrow d(i + 1, j + 1, 1) + 9 \times d(i, j, 1) \ (0 \leq j \leq 2) \]
となります。この時の増加分は \( 9 \times d(i, j, 1)\) ですが、これは、\( 0 \) 以外の数字の個数が \( j \) 個である状態から、\( i + 1\) 桁目を \( 0 \) 以外にすると、\( j + 1 \) 個の \( 0 \) 以外の数字を持つ状態に遷移すると考えると分かりやすいと思います。
\( n_i \neq 0\) のとき
\( d(i + 1, j, 0) \) の更新は、
\[ d(i + 1, j + 1, 0) \leftarrow d(i, j, 0) \ (0 \leq j \leq 2) \]
となります。
\( d(i + 1, j, 1) \) の更新は、\( i + 1 \) 桁目を \( 0 \) にする場合を考え、
\[ d(i + 1, j, 1) \leftarrow d(i, j, 1) \ (0 \leq j \leq 3) \]
とします。また、\( 0 \) 以外の値にすることを考えると、
\[ d(i + 1, j + 1, 1) \leftarrow d(i + 1, j + 1, 1) + 9 \times d(i, j, 1) + (n_{i + 1} – 1) \times d(i, j, 0) \ (0 \leq j \leq 2)\]
となります。\( (n_{i + 1} – 1) \times d(i, j, 0) \) の部分は \( N \) の \( i \) 桁目まで一致する数字から、\( i + 1 \) 桁目を \( 1 \leq x \leq n_{i + 1} – 1\) を満たす \( x \) にするときの遷移です。
したがって答えは、\( N \) の長さを \( l \) としたとき、\( d(l, K, 0) + d(l, K, 1) \) となります。
コード
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
string N;
int K, l;
ll d[105][4][2]{};
int main() {
cin >> N >> K;
l = N.length();
d[1][0][1] = 1;
d[1][1][0] = 1;
d[1][1][1] = N[0] - '0' - 1;
for (int i = 1; i < l; i++) {
ll s = N[i] - '0';
if (s == 0) {
for (int j = 0; j <= 3; j++) {
d[i + 1][j][0] = d[i][j][0];
}
for (int j = 0; j < 4; j++) {
d[i + 1][j][1] = d[i][j][1];
}
for (int j = 0; j < 3; j++) {
d[i + 1][j + 1][1] += d[i][j][1] * 9ll;
}
} else {
for (int j = 0; j < 3; j++) {
d[i + 1][j + 1][0] = d[i][j][0];
}
for (int j = 0; j < 4; j++) {
d[i + 1][j][1] = d[i][j][1] + d[i][j][0];
}
for (int j = 0; j < 3; j++) {
d[i + 1][j + 1][1] += d[i][j][1] * 9ll + d[i][j][0] * (N[i] - '0' - 1);
}
}
}
cout << d[l][K][0] + d[l][K][1] << "\n";
return 0;
}