Rustコトハジメ

プログラミング言語Rustと競プロに関する情報をお届けします。

【Educational DP-M Candies】累積和を使って計算量を削減するDP

M - Candies

問題

K個のあめをN人の子供でわけます。ただし、i番目の子供は0個以上ai以下しかとることが出来ません。

場合の数を求めよ。

制約: K=105, ai=105, N=100

方針

0番目がb個とると、残りはK-bをN-1人でわけろという再帰構造になっていることに着目します。

107かな?と思いきや、しっかり考えると、

  • N人
  • K個
  • ai個

の三重ループで、合計でO(NK2)=1012となります。ここから計算量を落とします。

f:id:akiradeveloper529:20190924205313j:plain

aiのループのところの絵を描いてみると、計算する値が、前の値の連続した領域の和をとっているだけなことがわかります。この累積和は、Kのループに入る前に計算出来て、それ以降使い回せます。

従って、オーダーはO(NK2)からO(NK)まで落ちることになります。

実装

もはや2次元配列も確保する必要もなく、空間量はO(K)で済みます。

fn solve() {
    input!{
        N:usize,K:usize,
        a:[usize;N]
    }
    let mut cur = vec![0;K+1];
    cur[K]=1;

    for i in 0..N {
        let mut cur_cumsum = cumsum1(&cur);
        let ai = a[i];
        for j in 0..K+1 {
            cur[j] = cur_cumsum[min(j+a[i], K)+1] + mo - cur_cumsum[j];
            cur[j] %= mo;
        }
    }
    println!("{}", cur[0]);
}

考察

ループの特徴に注目して計算量を削減するのがうまいところ。