貰うdpを累積和で高速化【競技プログラミング】

はじめに

• 貰うdpについて、前回の投稿で紹介したけど、今回は、貰うdpを復習するとともに、累積和との合わせ技を紹介したい。

問題例

• 教育的dpのm問題
  • N人の子供にK個の飴を与える方法は何種類有るか？(mod 1000000007で解答)
  • なお、子供iへは最大Ai個までしか与えることが出来ない。
    • つまりΣAi > K のときは解なし(0種類)となる。
• 入力例

3 4 // 3人の子供　4個の飴
1 2 3 // 子供[0]へは最大1個、子供[1]へは最大2個、子供[2]へは最大3個

貰うdp

• dp[i][j]を子供(0..<i)に対してj個の飴を与える場合の解とする。
• 貰うdpでの解答

extension Array { //二次元配列の表記を簡単にするextension
    init(_ cnt: Int, _ v: Element) {
        self = [Element](repeating: v, count: cnt)
    }
    init<T>(_ cnt2:Int, _ cnt1:Int, _ v:T) where Element == [T] {
        self = [[T]](cnt2,[T](cnt1,v))
    }
}

let (N,K) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1])}[0]
let A = readLine()!.split(separator:" ").map{Int($0)!}

main()

func main(){
    if A.reduce(0,+) < K { print(0) ; return } // 解がないときは0を返して終わり

    // dp配列
    var dp = [[Int]](N+1,K+1,0) // 初期値は0

    // 初期値
    dp[0][0] = 1 // 0人に0個を配る方法を1とする
    // どうして1なの？と根源的なことを聞かれると困っちゃうけど、
    // 後の推移でこれが種になるから、そういうものだと考えて。
    // まあ、5の0乗が1となっている様なものだと思ってください。悩んだら負け。

    // 貰うdpの遷移
    for n in 0..<N {
        for k in 0...K {
            for i in 0...min(A[n],k) {   // 子供[n]は飴をA[n]個までしか貰えない
                dp[n+1][k] += dp[n][k-i] // 子供[n]が0...iの飴を得る場合の数を合計
                                         // 子供[n]が0個の飴を得る場合の数はdp[n][k]
                                         // 子供[n]が1個の飴を得る場合の数はdp[n][k-1]
            }
        }
    }
    print(dp[N][K])
}

• 解説はコード中にコメントしたから、省略するね。
• ちなみに、配るdpなら、遷移部分のコードは次のとおり

    // 配るdpの遷移
    for n in 0..<N {
        for k in 0...K {
            for i in 0...min(A[n],K-k) { // 貰うdpからの変更部分
                dp[n+1][k+i] += dp[n][k] // 貰うdpからの変更部分
            }
        }
    }

• 貰うdpと配るdpの差というのは、端的に言えば、dpの遷移を表す式のどちらに一番内側のforループ変数(i)を含むか？と言うこと
  • つまり、iが左辺の式に含まれれば「配るdp」で、右辺の式に含まれれば「貰うdp」となる。
• 貰うdp、配るdpとも、提出するとTLEとなる。
• 制約を確認すると、$N<100、K<10^5、A[i]<K$となっていて、コードは3重forループとなっているから、$O(N･K^2) = O(10^{12})$なので、全然駄目ですね。

貰うdpに累積和を導入する

• 貰うdpの遷移のコードを見てみると...

    // 貰うdpの遷移
    for n in 0..<N {
        for k in 0...K {
            for i in 0...min(A[n],k) {
                dp[n+1][k] += dp[n][k-i] // 遷移式
            }
        }
    }

• for nループでdp[n][k] + dp[n][k-1] + dp[n][k-2] + ...が何度も繰り返し計算されていることが分かる。
• よって、コレを効率化することを考える。まぁ、どう考えても累積和だよね。タイトルにも累積和って書いてるしね。
• 具体的にはこんな感じ。
  • sumdpn[k] = Σdp[n][i] {i in 0..<k}
• コレを使って貰うdpの遷移のコードを書き換えると

    // 貰うdpの遷移
    for n in 0..<N {
    
        // ①累積和を作成
        var sumdpn = [Int](K+2,0) // 0...K+1の範囲となるから、K+2
        for k in 0...K {
            sumdpn[k+1] += sumdpn[k] + dp[n][k]
        }

        // ②貰うdpを累積和で高速化（forループを減らした）
        for k in 0...K {
            // for i in 0...min(A[n],k) {
            //     dp[n+1][k] += dp[n][k-i] // 遷移式
            // }
            dp[n+1][k] = (sumdpn[k+1] - sumdpn[k-min(A[n],k)]) // 遷移式
        }
    }

• はい、出来上がり！あれ？説明が足りない？
• まぁ、そうよね。説明します。
  • ①の累積和作成は、まぁ良いよね。分からない人は、過去の投稿を参考にしてね。
  • ②が頭こんがらがるから、整理します。
    • 累積和導入前の元の形は、dp[n][k]+dp[n][k-1]+...+dp[n][k-min(A[n],k)]だったね。
    • ここで、sumdpn[k+1]は、dp[n][k]+dp[n][k-1]+...+dp[n][0]
    • sumdpn[k-min(A[n],k)]は、dp[(k-1)-min(A[n],k-1)]+...+dp[n][0]
    • 以上です！お分かりいただけただろうか？
    • 分かりづらいポイントは、sudpn[k]のkとdp[n][k]のkの意味を1ズレて認識しやすいって事かな。
• よし、提出しよう！... あれ？RE(実行時エラー)になったよ。
  • mod 1000000007で解答するのを忘れてただけでした。遷移式を修正します。
  • dp[n+1][k] = (sumdpn[k+1] - sumdpn[k]) % 1000000007
• これで再提出すると100ms程度でACでした。

解答コード全体

extension Array {
    init(_ cnt: Int, _ v: Element) {
        self = [Element](repeating: v, count: cnt)
    }
    init<T>(_ cnt2:Int, _ cnt1:Int, _ v:T) where Element == [T] {
        self = [[T]](cnt2,[T](cnt1,v))
    }
}

let (N,K) = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1])}[0]
let A = readLine()!.split(separator:" ").map{Int($0)!}

main()

func main(){
    if A.reduce(0,+) < K { print(0) ; return }
    
    var dp = [[Int]](N+1,K+1,0)
    dp[0][0] = 1

    // 貰うdpの遷移
    for n in 0..<N {
    
        // 累積和を作成
        var sumdpn = [Int](K+2,0)
        for k in 0...K {
            sumdpn[k+1] += sumdpn[k] + dp[n][k]
        }

        for k in 0...K {
            dp[n+1][k] = (sumdpn[k+1] - sumdpn[k-min(A[n],k)]) % 1000000007
        }
    }
    
    print(dp[N][K])
}

• 計算量を確認すると、3重forループが2重forループになったお陰で、O(N･K)になったから、$O(10^7)$でTLEの問題はなくなったね。

おわりに

• 貰うdpの解法が少し体に馴染んだ感じがする。体に馴染ませるためには、修行あるのみ！