【桁DP】ABC154 - E - Almost Everywhere Zero

igasu_7

2024年3月13日 22:37

未履修のアルゴリズムについてアウトプット。

ABC154 E問題: Almost Everywhere Zero (diff 1298)

提出したコード
https://atcoder.jp/contests/abc154/submissions/50970179

整数$${N, K}$$が与えられる。
10進法で 0 でない数字が$${K}$$個現れる$${N}$$以下の正整数の個数を出力してね。

ゴリ押し

$${1 \leq K \leq3}$$なのと、Pythonの int に上限がないことを利用して無理やり通すこともできます。

初見の提出コード
https://atcoder.jp/contests/abc154/submissions/50947170

$${N}$$が3桁以下なら$${N}$$個の正整数を全探索する。
$${N}$$が4桁以上ならどの桁を 0 でない数字にするか$${K}$$ごとに場合分けして考える。

この機会に桁DPを勉強します。

桁DPとは

数値の桁に着目したDPです。

「ある条件」を満たす整数を全探索したいときに、ある桁の数字をひとつひとつ決めていくことで効率よく答えを求めるイメージ。

桁DPを学ぶにあたり、左からいくつかの桁の数値を固定したときのあり得る整数の個数について頭に入れておくと理解の助けになります。

桁DPの詳細な解説は他の方の記事にお任せします。

簡単なことから考える

「725以下の非負整数の個数を求めてください。」という問題を考えます。
明らかに726個ですが、桁DPの問題として解いてみます。

ある桁の数字を決めるにあたって、その数字を使うことで「725以下」という条件を満たさなくなる可能性を考慮しないといけません。
百の位での 7、十の位での 2、一の位での 5 がキーポイントになります。

左から$${i}$$桁目まで決まっていて$${i + 1}$$桁目を考えるとき、3つの遷移があることが分かります。

$${N}$$以下であることが確定するのは、$${N}$$の$${i + 1}$$桁目より小さい数を選択したときです。

今回の問題で考えること

上で見たように、桁DPでは$${i}$$桁目の時点で$${N}$$以下であることが確定しているかどうかの情報を持つ必要があります。

桁DPの基本形は以下になります。

$$
dp[i][smaller]: i桁目の時点でsmallerな整数の個数 \\
smaller = \begin{cases}
0 : N以下であることが未確定 \\
1 : N以下であることが確定 \\
\end{cases}
$$

これに加えて今回の問題では「 0 でない数字をいくつ使ったか」の情報も必要で、三次元DPをすることになります。

なお、0 は 0 でない数字が 0 個なので答えにはカウントされません。

三次元DP

$$
dp[i][smaller][k]: i桁目の時点で0でない数字をk個使っているsmallerな整数の個数 \\
smaller = \begin{cases}
0 : N以下であることが未確定 \\
1 : N以下であることが確定 \\
\end{cases}
$$

桁の番号は左から数えているため 0 桁目時点では$${N}$$以下であることが確定していません。よって$${dp[0][0][0] = 1}$$です。

$${i}$$桁目時点で$${N}$$以下であることが確定している、すなわち$${smaller = 1}$$であるとき$${i + 1}$$桁目の数字はなんでもOKです。

$${i}$$桁目時点で$${N}$$以下であることが確定していない、すなわち$${smaller = 0}$$であるとき、

$${i + 1}$$桁目の数字が$${N}$$の$${i + 1}$$桁目と等しい場合は、まだ$${N}$$以下であることが確定しないままです。$${(smaller = 0)}$$
$${i + 1}$$桁目の数字が$${N}$$の$${i + 1}$$桁目より小さい場合は、$${N}$$以下であることが確定します。$${(smaller = 1)}$$
$${N}$$の$${i + 1}$$桁目より大きい数字は選べません。

N = int(input())
K = int(input())
S = [i for i in map(int, str(N))]
M = len(S)
""" 
dp[i][smaller][k]: i 桁目までに 0 でない数字を k 個使っている整数の個数
(smaller = 0: i 桁目の時点で N 以下であることが未確定) 
(smaller = 1: i 桁目の時点で N 以下であることが確定)
"""
dp = [[[0] * (K + 1) for _ in range(2)] for _ in range(M + 1)]
dp[0][0][0] = 1
for i in range(M):
    for k in range(K + 1):
        for j in range(10):
            if k == K and j > 0:
                break
            dk = (j != 0)
            dp[i + 1][1][k + dk] += dp[i][1][k]
            if S[i] == j:
                dp[i + 1][0][k + dk] += dp[i][0][k]
            elif S[i] > j:
                dp[i + 1][1][k + dk] += dp[i][0][k]
ans = dp[M][0][K] + dp[M][1][K]
print(ans)

二次元DP

$${smaller}$$を別々のリストとして分割することでDPリストの次元を減らすことができます。

$$
dp[i][k]: i桁目の時点で0でない数字をk個使っている整数の個数 \\

dp0 : N以下であることが未確定 \\
dp1 : N以下であることが確定 \\
$$

N = int(input())
K = int(input())
S = [i for i in map(int, str(N))]
M = len(S)
""" 
dp[i][k]: i 桁目までに 0 でない数字を k 個使っている整数の個数
(dp0: i 桁目の時点で N 以下であることが未確定) 
(dp1: i 桁目の時点で N 以下であることが確定)
"""
dp0 = [[0] * (K + 1) for _ in range(M + 1)]
dp1 = [[0] * (K + 1) for _ in range(M + 1)]
dp0[0][0] = 1
for i in range(M):
    for k in range(K + 1):
        for j in range(10):
            if k == K and j > 0:
                break
            dk = (j != 0)
            dp1[i + 1][k + dk] += dp1[i][k]
            if S[i] == j:
                dp0[i + 1][k + dk] += dp0[i][k]
            elif S[i] > j:
                dp1[i + 1][k + dk] += dp0[i][k]
ans = dp0[M][K] + dp1[M][K]
print(ans)

一次元DP

DPの遷移には1つ前の桁についての情報だけあればよいので、さらに分割して長さ$${K + 1}$$の一次元リスト4つで解くこともできます。

$$
dp[k]: i桁目の時点で0でない数字をk個使っている整数の個数 \\

dp0, ndp0 : N以下であることが未確定 \\
dp1, ndp1 : N以下であることが確定 \\
$$

N = int(input())
K = int(input())
S = [i for i in map(int, str(N))]
M = len(S)
""" 
dp[k]: i 桁目までに 0 でない数字を k 個使っている整数の個数
(dp0, ndp0: i 桁目の時点で N 以下であることが未確定) 
(dp1, ndp1: i 桁目の時点で N 以下であることが確定)
"""
dp0 = [0] * (K + 1)
dp1 = [0] * (K + 1)
dp0[0] = 1
for i in range(M):
    ndp0 = [0] * (K + 1)
    ndp1 = [0] * (K + 1)
    for k in range(K + 1):
        for j in range(10):
            if k == K and j > 0:
                break
            dk = (j != 0)
            ndp1[k + dk] += dp1[k]
            if S[i] == j:
                ndp0[k + dk] += dp0[k]
            elif S[i] > j:
                ndp1[k + dk] += dp0[k]
    dp0 = ndp0
    dp1 = ndp1
ans = dp0[K] + dp1[K]
print(ans)

おまけ

forループを1つ減らせます。

DPリストを長さ$${K + 2}$$にして if 文を減らしています。

N = int(input())
K = int(input())
S = [i for i in map(int, str(N))]
M = len(S)
""" 
dp[k]: i 桁目までに 0 でない数字を k 個使っている整数の個数
(dp0, ndp0: i 桁目の時点で N 以下であることが未確定) 
(dp1, ndp1: i 桁目の時点で N 以下であることが確定)
"""
dp0 = [0] * (K + 2)
dp1 = [0] * (K + 2)
dp0[0] = 1
for i in range(M):
    ndp0 = [0] * (K + 2)
    ndp1 = [0] * (K + 2)
    for k in range(K + 1):
        if S[i] == 0:
            ndp0[k] += dp0[k]
            ndp1[k] += dp1[k]
        else:
            ndp0[k + 1] += dp0[k]
            ndp1[k] += dp0[k] + dp1[k]
        ndp1[k + 1] += dp1[k] * 9 + dp0[k] * max(0, S[i] - 1)
    dp0 = ndp0
    dp1 = ndp1
ans = dp0[K] + dp1[K]
print(ans)

桁DPに基本形があるとは知りませんでした。

今回の問題は「$${N}$$以下であることが確定している」ことが分かりやすかったですが、そうでない問題もあるようです。

類題: ABC117 - D - XXOR (diff: 1423)

提出コード: https://atcoder.jp/contests/abc117/submissions/50998292

この記事が気に入ったらサポートをしてみませんか？