はまやんはまやんはまやん

hamayanhamayan's blog

Candies [Educational DP Contest / DP まとめコンテスト M]

https://atcoder.jp/contests/dp/tasks/dp_m

解説

https://atcoder.jp/contests/dp/submissions/3948814

dp[i][k] := i人目の子供まで配り終えて、残りがk個である組み合わせ
dp[i][k]からの遷移を考えると、i+1番目の子供は0~A[i]個取ることができる。
(A[i+1]が本当は正しいが、Aの添字を0-indexedと考える)

rep(i, 0, N) rep(k, 0, K+1) rep(nxt, 0, A[i] + 1) dp[i+1][k-nxt] += dp[i][k];

これはO(NK^2)で間に合わない。
このDPはある状態から他の状態に遷移していて、組み合わせを配っているように見えるため『配るDP』と呼ばれている。
 
よくあるテクとして、配るDPを『貰うDP』に帰ることで高速化の余地が出てくる場合がある。
『貰うDP』とは、状態を更新する側を主観として、その状態への遷移元の組み合わせをまとめる方法である。
具体的にはdp[i][k]はdp[i-1][k], dp[i-1][k+1], dp[i-1][k+2], ..., dp[i-1][k+A[i-1]]から遷移される。
よって、dp[i][k] = dp[i-1][k] + dp[i-1][k+1] + dp[i-1][k+2] + ... + dp[i-1][k+A[i-1]]となる。
これはBIT(または累積和)を使えば高速化できそう。
BITを使ってDPを構築すれば高速に取得できる。

int N, K, A[101];
BIT<mint> dp[101];
//---------------------------------------------------------------------------------------------------
void _main() {
    cin >> N >> K;
    rep(i, 0, N) cin >> A[i];
 
    rep(i, 0, N + 1) dp[i].init(K + 1);
    dp[0].add(K, 1);
 
    rep(i, 0, N) rep(k, 0, K + 1) {
        dp[i+1].add(k, dp[i].get(k, min(K, k+A[i])+1));
    }
 
    cout << dp[N].get(0,1) << endl;
}