【Atcoder】M - Candies【DPまとめコンテスト】

問題

結論

dp[i][j] := i人目まででj個ぴったり配る場合の数

def getInts():
    return map(int, input().split())


def getIntsArray():
    return [int(x) for x in input().split()]


N, K = getInts()
Capacities = getIntsArray()
MOD = pow(10, 9) + 7

dp = [[0] * (K + 1) for _ in range(N + 1)]
dp[0][0] = 1
for i in range(1, N + 1):
    cumsum = [0] * (K + 1)
    cumsum[0] = dp[i - 1][0]
    for k in range(1, K + 1):
        cumsum[k] = cumsum[k - 1] + dp[i - 1][k]
        cumsum[k] %= MOD
    for k in range(K + 1):
        capacity = Capacities[i - 1]
        min_limit = k - capacity - 1
        if min_limit >= 0:
            dp[i][k] = (cumsum[k] - cumsum[min_limit]) % MOD
        else:
            dp[i][k] = cumsum[k]
print(dp[N][K])

思考順序

愚直な実装

素直に実装すると

dp = [[0] * (K + 1) for _ in range(N + 1)]
dp[0][0] = 1
for i in range(1, N + 1):
    for k in range(K + 1):
        capacity = Capacities[i - 1]
        for n in range(capacity + 1):
            if k - n >= 0:
                dp[i][k] += dp[i - 1][k - n]
                dp[i][k] %= MOD
print(dp[N][K])

これだと計算量が$O(NK^2)$

ループを減らす

例で考えてみましょう

ex) N = 3, K = 4, a = (1,2,3)
最終的なDPテーブルは以下の通り

[1, 0, 0, 0, 0]
[1, 1, 0, 0, 0]
[1, 2, 2, 1, 0]
[1, 3, 5, 6, 5]

例えばdp[2][2]を求めるとき、考える遷移は

$$
dp[2][2] = dp[1][0] + dp[1][1] + dp[1][2]
$$

つまり、1つ前の行の累積和$cumsum$が計算できていれば、$cumsum[a_i]$で求められます

区間の和は累積和の差

また、別の例を考えてみます

ex) N=3, K=4, a = (2, 1, 1)

dp[2][2]を求めるとき、考える遷移は

$$
dp[2][2] = dp[1][1] + dp[1][2]
$$

です。

これもdp1行目の累積和$cumsum$が計算できていれば

$cumsum[2]-cumsum[2-a_2 - 1] = cumsum[2] - cumsum[0]$

で求めることができます。

• 1列目から2列目までの和 = 2列目までの累積和 - 0列目までの累積和

実装

1. 1行前の累積和を計算する
2. 累積和を使って、dpテーブルを更新する

こうすると、ループを1つ減らせて計算量は$O(NK)$

for i in range(1, N + 1):
    cumsum = [0] * (K + 1)
    cumsum[0] = dp[i - 1][0]
    for k in range(1, K + 1):
        cumsum[k] = cumsum[k - 1] + dp[i - 1][k]
        cumsum[k] %= MOD
    for k in range(K + 1):
        capacity = Capacities[i - 1]
        min_limit = k - capacity - 1
        if min_limit >= 0:
            dp[i][k] = (cumsum[k] - cumsum[min_limit]) % MOD
        else:
            dp[i][k] = cumsum[k]

MODで割るのを忘れずに(1敗)