PythonでABC154-E(桁DP)

問題

ABC154-Eを理解して見せる！

総数

まず、桁DPを用いてN以下の数の総数を求める。
314159について考えてみる。
例えば、313まで決まっていると、それより下はどんな数字でも許される。
一方、314という決まり方をすると、1個下は0か1しか許されない。
すなわち、どの桁までを決めるのかと、その決め方でN未満が確定するのかという情報が必要だ。
dp[i][j]: i桁目までを決めて、N未満で確定ならj=1、Nと一致しているならばj=0
と定める。
初期値は、0桁目までを使ってこの時点ではNと一致しているので
dp[0][0] = 1
とする。
その後、桁数を増やしていくが、次の桁について、Nと一致している個数は常に1(1つ前と同じ)、N未満の個数は、1つ前のN未満確定のものに0~9の10をかけたものと1つ前のNに等しいものに0~その桁の数-1をかけたものになるので、漸化式は以下。
dp[i][0] = dp[i - 1][0]
dp[i][1] = dp[i - 1][1] * 10 + dp[i - 1][0] * int(str(N)[i - 1])

コード

.py

N = input()
m = len(N)
dp = [[0] * 2 for _ in range(m + 1)]
dp[0][0] = 1

for i in range(1, m + 1):
    dp[i][0] = dp[i - 1][0]
    dp[i][1] = dp[i - 1][1] * 10 + dp[i - 1][0] * int(N[i - 1])

ans = dp[m][0] + dp[m][1]

print(ans)

0でない数をK個含むN以下の数の総数

dp[i][j][k]: 0でない数をk個含むという条件追加
i桁目の数をlとして、
dp[i][0][k] = dp[i - 1][0][k] * (l == 0),dp[i - 1][0][k - 1] * (l != 0)
dp[i][1][k] = dp[i - 1][1][k] + dp[i - 1][1][k - 1] * 9 + dp[i - 1][0][k] * 1 * (l != 0) + dp[i - 1][0][k - 1] * (l - 1) * (l != 0)
になると思うんですが無駄がありそうなので教えてください><

考え方

dp[i][0][k]:lが0だった場合には与えられた数と一致させるにはkはそのままになるのでdp[i -1][0][k]となり、そうでないときには0でない数が1つ増えるのでdp[i - 1][0][k - 1]。

dp[i][1][k]:まず、dp[i - 1][1][k]から次の桁が0でないもの1つ、dp[i - 1][1][k - 1]から次の桁が1~9の9つが遷移してくる。次に、lが0でなければ、次の桁を0としたdp[i - 1][0][k]が、また次の桁が1~(l - 1)となるdp[i - 1][0][k - 1]がl - 1個遷移してくる。
(わかりづら杉)

コード

.py

N = input()
K = int(input())
m = len(N)
dp = [[[0] * (K + 1) for _ in range(2)] for _ in range(m + 1)]
dp[0][0][0] = 1

for i in range(1, m + 1):
    l = int(N[i - 1])
    for k in range(K + 1):
        dp[i][1][k] += dp[i - 1][1][k]
        if l != 0:
            dp[i][1][k] += dp[i - 1][0][k]
        else:
            dp[i][0][k] += dp[i - 1][0][k]
        if k - 1 >= 0:
            dp[i][1][k] += 9 * dp[i - 1][1][k - 1]
            if l != 0:
                dp[i][0][k] += dp[i - 1][0][k - 1]
                dp[i][1][k] += (l - 1) * dp[i - 1][0][k - 1]

print(dp[m][0][K] + dp[m][1][K])

あとがき

アドバイスお願いします。。。。
質問も答えられる範囲で答えるつもりですが。。。。

参考文献

桁DP入門