2023-04-30

EDPC (Educational DP Contest) M - Candies

AtCoder EDPCとTDPC DP 累積和 DP高速化 DP高速化:累積和多項式・FPS(形式的冪級数) 制約条件:総和=K 制約条件:和や差が一定値数え上げ問題制約条件:xi<=ai FPSテク:漸化式の母関数を求めるそのまま覚えたい典型問題【問題集】DPのステップアップナップサックDP

DP 高速化に累積和を使う問題！

問題へのリンク

問題概要

$N$ 人の子ども $1, 2, \dots, N$ に、 $K$ 個の飴を分けることにした。

ただし、子ども $i$ に分ける飴は、 $0$ 個以上 $a_{i}$ 個以下のする必要がある。

各子どもへの飴の分け方の総数を、1000000007 で割ったあまりを求めよ。

制約

$1 \le N \le 100$
$0 \le K \le 10^{5}$
$0 \le a_{i} \le K$

解法：まずは DP を組む

DP に慣れて来ると、次のような配列 dp を定義することで解けそうだと思えると思われる。

dp[i][j] ← 子ども $0, 1, \dots, i-1$ に対して、 $j$ 個の飴を分ける方法の総数

そうすると、DP の遷移は次のように作れる。なお、DP 遷移を作るときは「配る DP」と「集める DP」のどちらでもよいのだが、DP 高速化を見据える場合は、集める DP にするといいケースが多いように思う。

集める DP で考えると、次のようになる。

dp[i + 1][j] = dp[i][j] + dp[i][j - 1] + ... + dp[i][j - min(j, a[i])]

ここで、右辺において

dp[i][j] は「 $i$ 人目の子どもに飴を 0 個あげる」場合
dp[i][j - 1] は「 $i$ 人目の子どもに飴を 1 個あげる」場合
...
dp[i][j - min(j, a[i])] は「 $i$ 人目の子どもに飴を $\min(j, a_{i})$ 個あげる」場合

を表している。

この漸化式に従って DP を実装すると、dp[i + 1][j] を計算するのに要する計算量が $O(K)$ であることから、全体の計算量は $O(NK^{2})$ となる。

このままでは TLE となるので、高速化する必要がある。

累積和で DP 高速化

DP の遷移式

dp[i + 1][j] = dp[i][j] + dp[i][j - 1] + ... + dp[i][j - min(j, a[i])]

の右辺をよく見ると、累積和をつかって高速化できそうな形になっていることがわかる。

具体的には、

dp[i][0], dp[i][1], dp[i][2], ...

という数列を考える。このとき、上式の右辺は、この数列の連続する要素の和となっている。そこで、この数列の累積和を考えよう。累積和を sum_dp とおくと、

sum_dp[0] = 0
sum_dp[1] = dp[i][0]
sum_dp[2] = dp[i][0] + dp[i][1]
sum_dp[3] = dp[i][0] + dp[i][1] + dp[i][2]
...
sum_dp[K+1] = dp[i][0] + dp[i][1] + dp[i][2] + ... + dp[i][K]

というようになる。このとき、DP の遷移式は次式のように表せる！！

dp[i + 1][j] = sum_dp[j + 1] - sum_dp[j - min(j, a[i])]

こうすると、計算量が $O(NK)$ に改善される。より詳細には

各に対して
- dp[i] の累積和 sum_dp を作る部分：計算量 $O(K)$
- sum_dp を用いて、各 $j$ に対して dp[i+1][j] を求める部分：計算量 $O(K)$

となるので、合計で $O(NK)$ となる。

解答例

ACL の modint を使った。

#include <bits/stdc++.h>
using namespace std;

#include "atcoder/modint.hpp"
using mint = atcoder::modint1000000007;

int main() {
    int N, K;
    cin >> N >> K;
    vector<int> a(N);
    for (int i = 0; i < N; ++i) cin >> a[i];
    
    vector<vector<mint>> dp(N+1, vector<mint>(K+1, 0));
    dp[0][0] = 1;
    for (int i = 0; i < N; ++i) {
        // dp[i] の累積和を求める
        vector<mint> sum_dp(K+2, 0);
        for (int j = 0; j <= K; ++j) {
            sum_dp[j+1] = sum_dp[j] + dp[i][j];
        }
        
        // dp[i+1] を求める
        for (int j = 0; j <= K; ++j) {
            dp[i+1][j] = sum_dp[j+1] - sum_dp[j-min(j, a[i])];
        }
    }
    cout << dp[N][K].val() << endl;
}