AtCoder Beginner Contest 154

E - Almost Everywhere Zero

atcoder.jp

解法

AtCoderの解説に準じた解法で解きます。桁DPとよばれるタイプのDPを用います。
入力例3に対してdp_tableを埋めていく指針を説明します。

[入力例3]
N = 314159
K = 2

このケースのdp_tableを先に描くと以下のようなものになります。

添え字が3つもあるので考えるのがちょっと大変です。
具体的に、例えば表中の緑色のセルが意味するところは、
「4桁目の時点で、0でない数字が2つあり、4桁目までの数字が314159と一致していない数字の個数はX個」
です。以下のような数字が具体例として挙げられます。

1100**, 3009**, 0890**, 0013**, ... は全部でX個

別例として、上の表中の緑色のセルが意味するところは、
「3桁目の時点で、0でない数字が1つあり、3桁目までの数字が314159と一致している数字の個数はY個」
です。
「3桁目までの数字が314159と一致している」数字の候補は

314***

のただ1つです。これは既に「0でない数字が1つ」ではないので、Yは明らかに0です。

以下表中のA,Bを足したものが本問題の解答になります。 Aは前例同様0です。なぜなら、「6桁目までの数字が314159と一致している」数字の候補は314159ただ1つであり、これは「0でない数字が2つ」ではないからです。
Bは「6桁目までの数字が314159と一致していない」数字であり、候補は314158個あります。
最終的に「314159」と「314159未満」の数字の中から条件に合致するものの個数を算出し、出力するという流れになります。

解説

初期状態

桁dpの初期条件は以下のような配置になります。 dp[0][0][0]、つまり「0桁目の時点で0でない数字が0つあり、0桁目までの数字が314159と一致している数字の個数」、は1個、それ以外は0個です。
解説動画の1:10:45あたりでこの初期状態に関する解説があるので詳しい話はそちらに譲ります。

1桁目への遷移

i = 1の行を考えます。このテストケースで1桁目にはどのような数字がとれるでしょうか。
1桁目の場合、明らかに0~3のいずれかになります。その中でどの数字を取ったかで遷移先が変わります。 1. 0をとった場合
取った数字が0なので、0でない数字は遷移元と同じ個数のj=0のままです。
1桁目に0をとり、1桁目の数字が314159と一致しなくなったので、k=1に遷移します。
遷移元のパターンが1通りで、新しく選ぶ数字が[0]の1つなので、遷移先のパターンは1個(=1 * 1)増加します。

2. 1,2を取った場合
取った数字が0ではないので、0でない数字は遷移元から1つ増えてj=1になります。
1桁目に1もしくは2をとり、1桁目の数字が314159と一致しなくなったので、k=1に遷移します。
遷移元のパターンが1通りで、新しく選ぶ数字が[1,2]の2つなので、遷移先のパターンは1個(=1 * 2)増加します。

3. 3を取った場合
取った数字が0ではないので、0でない数字は遷移元から1つ増えてj=1になります。
1桁目に3をとり、1桁目の数字は314159と一致しているので、k=0のままです。
遷移元のパターンが1通りで、新しく選ぶ数字が[3]の1つなので、遷移先のパターンは1個(=1 * 1)増加します。

2桁目への遷移

i = 2の行を考えます。このテストケースで2桁目にはどのような数字がとれるでしょうか。
これは1桁目の数字に依存します。もし1桁目が既に3、つまり、「1桁目までの数字が314159と一致している(k=0)」場合、2桁目に取れるのは0か1だけです。（例えばもし2をとってしまうと 32**** となり、その時点で314159をオーバーしてしまう） 逆に1桁目が3でない、つまり、「1桁目までの数字が314159と一致していない(k=1)」場合、2桁目には0~9の数字を自由にとることができます。
dp[1][0][1]からの遷移を考えます。この1個(0*****)に関しては、2桁目の数字を0~9の中から選ぶことができます。
0を取った場合は、0でない数字の個数は変わらず、j=0のままです。
1~9を取った場合は、0でない数字は1つ増え、j=1に遷移します。
どちらのケースにおいても、既にNと一致していない数字からの遷移であるため、k=1のままです。
遷移先の個数は、遷移元の個数 x 取れる数字の数だけ増加します。

dp[1][1][1]からの遷移を考えます。この2個(1*****, 2*****)に関しても、2桁目の数字を0~9の中から選ぶことができます。
0を取った場合は、0でない数字の個数は変わらず、j=1のままです。
1~9を取った場合は、0でない数字は1つ増え、j=2に遷移します。
どちらのケースにおいても、既にNと一致していない数字からの遷移であるため、k=1のままです。

dp[1][1][0]からの遷移を考えます。この1個(3*****)がとれる2桁目の数字は、0, 1のみです。
0を取った場合は、0でない数字の個数は変わらず、j=1のままです。
またこの場合、2桁目の段階で314159と一致しなくなったので、k=1になります
1を取った場合は、0でない数字は1つ増え、j=2に遷移します。
この場合、2桁目の段階で31****となり、314159と一致しているので、k=0のままです

遷移まとめ

遷移は以下のようになります。

x = "i+1桁目で取れる数の個数"
for s in range(x):
    dp[i+1][next_j][next_k] += dp[i][j][k]

ただし、

if k==0:
    x = (Nのi+1桁目の数)
else:
    x = 9
# k==0(i桁目までNと一致）なら、i+1桁目は自由にはとれない(Nのi+1桁目の数まで)

if s == 0:
    next_j = j 
else:
    next_j = j+1
#  次にとる数が0のときのみjはそのまま

if k==0 and s==(Nのi+1桁目の数):
    next_k = 0
else:
    next_k = 1
# k==0 (i桁目までNと一致)で、さらにi+1桁目もNのi+1桁目と一緒なら、next_k = 0

これを踏まえて3桁目への遷移を書くと以下のようになります。 Kをこえたjに関しては走査から除外して大丈夫です。この問題ではjが小さくなる方向に遷移することはないからです。

実装

N=int(input())
K=int(input())

def dp(N,K):
    keta = len(str(N))
    table = [[[0,0] for i in range(K+1)] for j in range(keta+1)]
    table[0][0][0]=1
    for i in range(keta):
        for j in range(K+1):
            for k in range(2):
                digit = int(str(N)[i]) if k==0 else 9
                for x in range(digit+1):
                    next_j = j if x == 0 else j + 1
                    next_k = 0 if k == 0 and x == digit else 1
                    if next_j > K: continue
                    table[i+1][next_j][next_k]+=table[i][j][k]
    return table[keta][K][0] + table[keta][K][1]
print(dp(N,K))

レート変化