題目的難度顏色使用 Luogu 上的分級，由簡單到困難分別為 🔴🟠🟡🟢🔵🟣⚫。

🔗 🟢 AT_dp_m Candies

Problem Statement

題目簡述

有 NN 個小孩要分配 KK 顆糖果，第 ii 個小孩最多拿 aia_i 顆。求恰好分完所有糖果的方案數，答案對 109+710^9+7 取模。

Constraints

約束條件

  • 1≤N≤1001 \le N \le 100
  • 0≤K≤1050 \le K \le 10^5
  • 0≤ai≤K0 \le a_i \le K

思路：DP + 前綴和優化

狀態定義

定義 f[i][k]f[i][k] 為前 ii 個小孩恰好分配 kk 顆糖的方案數。

初始狀態 f[0][0]=1f[0][0] = 1（0 個小孩分 0 顆糖的方案數為 1）。

狀態轉移

考慮第 ii 個小孩拿 j∈[0,ai]j \in [0, a_i] 顆糖，則：

f[i][k]=∑j=0min⁡(ai,k)f[i−1][k−j]f[i][k] = \sum_{j=0}^{\min(a_i, k)} f[i-1][k-j]

樸素實作需要 O(NK2)\mathcal{O}(NK^2)，會超時。

前綴和優化

核心觀察

轉移式中的累加項是 f[i−1]f[i-1] 在區間 [max⁡(0,k−ai), k][\max(0, k-a_i),\, k] 上的連續和。

令前綴和 S[k]=∑j=0k−1f[i−1][j]S[k] = \sum_{j=0}^{k-1} f[i-1][j]（使用 initial=0 使索引偏移 1），則：

f[i][k]=S[k+1]−S[max⁡(0,k−ai)]f[i][k] = S[k+1] - S[\max(0, k - a_i)]

這樣每次轉移只需 O(1)\mathcal{O}(1)，總體優化至 O(NK)\mathcal{O}(NK)。

滾動陣列

由於 fif_i 只依賴 fi−1f_{i-1}，可使用兩個一維陣列交替滾動，空間複雜度為 O(K)\mathcal{O}(K)。

複雜度分析

  • 時間複雜度：O(NK)\mathcal{O}(NK)
  • 空間複雜度：O(K)\mathcal{O}(K)

Code

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
from itertools import accumulate

MOD = int(1e9 + 7)

def solve():
    N, K = map(int, input().split())
    A = list(map(int, input().split()))
    assert len(A) == N

    f = [0] * (K + 1)
    f[0] = 1
    for x in A:
        nf = [0] * (K + 1)
        # for k in range(K + 1):
        #     for j in range(x + 1):
        #         if k - j >= 0:
        #             nf[k] = (nf[k] + f[k - j]) % MOD
        #         else:
        #             nf[k] = nf[k]
        s = list(accumulate(f, lambda x, y: (x + y) % MOD, initial=0))
        for k in range(K + 1):
            nf[k] = (s[k + 1] - s[max(0, k - x)]) % MOD
        f = nf
    print(f[K])

if __name__ == "__main__":
    solve()