トップ > AtCoder > EDPC - M - Candies

EDPC - M - Candies

AtCoder AtCoder-EDPC DP 和-累積和 DP-累積和DP いもす法 Python

ある区間から遷移する問題．
貰うDPなら累積和（ある区間からの遷移），配るDPならいもす法（ある区間への遷移）となる．

個人的には，いもす法の方がわかりやすい．

方法1 (貰うDP & 累積和)

貰うDPを考えると，「ある区間からの遷移」になる．

考え方

dp[i][j]=$i$番目までで合計$j$個配る方法の数とすれば良さそう．貰うDPを考えると

N, K = map(int, input().split())
A = list(map(int, input().split()))

mod = 10 ** 9 + 7
dp = [[0] * (K + 1) for _ in range(N + 1)]
dp[0][0] = 1

for i in range(N):
    for j in range(K + 1):
        for k in range(A[i] + 1):
            if j - k >= 0:
                dp[i + 1][j] += dp[i][j - k]
                dp[i + 1][j] %= mod

print(dp[N][K])

となり，ある区間から遷移になっている．式で書くと

\begin{aligned}
\mathrm{dp}[i+1][j]
&=\sum_{k=0}^{\min(j, A[i])} \mathrm{dp}[i][j - k] \\
&=\sum_{l=\max(0, j - A[i])}^{j} \mathrm{dp}[i][l]
\end{aligned}
だから，累積和が使えそう．

回答例

各$i$で累積和を求めておけば，$k$のループは不要になる．

N, K = map(int, input().split())
A = list(map(int, input().split()))

mod = 10 ** 9 + 7
dp = [[0] * (K + 1) for _ in range(N + 1)]
dp[0][0] = 1

for i in range(N):
    cumsum = [0] * (K + 2)
    for j in range(K + 1):
        cumsum[j + 1] += cumsum[j] + dp[i][j]
    for j in range(K + 1):
        dp[i + 1][j] = cumsum[j + 1] - cumsum[max(0, j - A[i])]
        dp[i + 1][j] %= mod

print(dp[N][K])

方法2 (配るDP & いもす法)

配るDPを考えると，「ある区間への遷移」になる．

考え方

dp[i][j]=$i$番目までで合計$j$個配る方法の数とすれば良さそう．配るDPを考えると

N, K = map(int, input().split())
A = list(map(int, input().split()))

mod = 10 ** 9 + 7
dp = [[0] * (K + 1) for _ in range(N + 1)]
dp[0][0] = 1

for i in range(N):
    for j in range(K + 1):
        for k in range(A[i] + 1):
            if j + k <= K:
                dp[i + 1][j + k] += dp[i][j]
                dp[i + 1][j + k] %= mod

print(dp[N][K])

となり，「ある区間への遷移」になっている．このタイプの計算は「いもす法」を使える．

回答例

遷移はimosで出入りだけ管理して，遷移がすべて終わったら累積和を計算してdpを求める．

N, K = map(int, input().split())
A = list(map(int, input().split()))
mod = 10 ** 9 + 7

dp = [[0] * (K + 1) for _ in range(N + 1)]
dp[0][0] = 1

for i in range(N):
    imos = [0] * (K + 1)
    for j in range(K + 1):
        imos[j] += dp[i][j]
        imos[j] %= mod
        if j + A[i] + 1 <= K:
            imos[j + A[i] + 1] -= dp[i][j]
            imos[j + A[i] + 1] %= mod
    dp[i + 1][0] = imos[0]
    for j in range(K):
        dp[i + 1][j + 1] = dp[i + 1][j] + imos[j + 1]
        dp[i + 1][j + 1] %= mod

print(dp[N][K])

よく考えてみれば，imosという配列を用意しなくても，dpをそのまま利用できる．

N, K = map(int, input().split())
A = list(map(int, input().split()))
mod = 10 ** 9 + 7

dp = [[0] * (K + 1) for _ in range(N + 1)]
dp[0][0] = 1

for i in range(N):
    for j in range(K + 1):
        dp[i + 1][j] += dp[i][j]
        dp[i + 1][j] %= mod
        if j + A[i] + 1 <= K:
            dp[i + 1][j + A[i] + 1] -= dp[i][j]
            dp[i + 1][j + A[i] + 1] %= mod
    for j in range(K):
        dp[i + 1][j + 1] += dp[i + 1][j]
        dp[i + 1][j + 1] %= mod

print(dp[N][K])

背景

E - RLEの類題として．
（ABC249E - RLE - 競プロはじめました）