挿入dpについて【競技プログラミング】

はじめに

• 教育的DP問題のt問題が挿入dpを使って解く問題のようだけど、初めて知ったアルゴリズムなので調べてみた。
• とりあえず、同じ挿入dpでも定番っぽい問題のほうから解いてみる。
• なお、挿入dpとは、「問題で与えられた初期配列」を新たに配列に並び替えていく際、まず、新たな配列用の空配列を用意し、初期配列の先頭から配列を取り出して、新たな配列に挿入していくときにdpを用いる手法のこと。まぁ、名前の通りだね。

定番の挿入dp問題

• サイズNの配列Aが用意されます。要素はダブることがあります。例えば、A = [5 1 4 4 2]の4がダブっている感じです。
• この配列を並べ替えてBとするとき、隣同士の要素B[i]<B[i+1]となるiの数が丁度K個となる並び替え方法はいくつあるか？(mod 998244353)

解き方

• まず、予め、Aをソートしておきます。最初の例であれば、A = [1 2 4 4 5]とします。
• そして、並び替え先の新配列として、B=[]を最初に考えて、Aの要素を新配列に挿入していくこととします。最初の挿入結果は必ず、[1]となり、2つ目の挿入結果は、[1 2]か[2 1]のいずれかとなります。
• Aの要素を追加して行きますが、挿入による隣同士が「>」となる要素のペア数の変化は、「1つ増加する」か「変わらない」のどちらかであり、「1つ減少したり」「2つ増加する」ことはないです。
• で、「1つ増加する」のは、挿入するAの要素をxとして、x>B[i]≧B[i+1]の関係となるB[i]とB[i+1]の間に挿入するときに限ります。理由は下記の通り。
  • Aはソート済みであるため、x ≧ max(B[i],B[i+1])となるので、B[i]<B[i+1]の間にxを挿入しても、「変わらない」となります。
  • x = B[i]の時にB[i]とB[i+1]の間に挿入しても「変わらない」となります。
• よって、組み合わせの数をカウントするのに都合の良いdpの設定は以下の通りです。
  • dp[n][k]:=ソート後のAの小さい要素からn個の要素を追加したとき、B[i]<B[i+1]となるiの数がk個となる組み合わせの数。なお、この時は当然、B[i]≧B[i+1]となるiの数はn-k-1個となる。
• 上記の考えに基づき、dpを使った遷移を書くと、以下の通り。

dp[0][0] = 1
for n in 0..<N {
    for k in 0...K {
        //「変わらない」パターン 
        dp[n+1][k] += dp[n][k] * (k + 1)

        //「1つ増加する」パターン
        dp[n+1][k+1] += dp[n][k] * (n - k)
    }
}

• 説明を簡単にするため、上記コードはAの要素の値は全て異なると考えた時の遷移となっています。
• 見てのとおり、「<」となるペアの数が「変わらない」パターンと、「1つ増加する」パターンに分かれます。
  • 「変わらない」パターンでは、B[i]<B[i+1]となるペア(k個)の間に差し込むけど、あれ？(K+1)を乗じてるのは何故？と思うよね。実は差し込まなくても、Bの先頭に置いても「変わらない」から、この分が1追加されて(k+1)となる。なお、Bの末尾に置くと、「1つ増加する」ので、末尾に置くパターンは考えない。
  • 「1つ増加する」パターンでは、B[i]≧B[i+1]となるペア(n-k-1個)の間に差し込むときだけでなく、Bの末尾に置いても「1つ増加する」から(n-k)を乗じます。
• 上記コードについては、お分かりいただけただろうか？
• つぎに、同じ数が続いたときを考慮したコードを書きます。
  • 新たな変数sameを追加し、初期状態は0とします。これは、同じ数が続いたときに1,2,3と増加し、違う数になったら0に戻ります。

dp[0][0] = 1
var same = 0
for n in 0..<N {
    if n > 0 && A[n-1] == A[n] {same += 1} else {same = 0} 
    for k in 0...K {
        //「変わらない」パターン     
        dp[n+1][k] += dp[n][k] * (k + 1 + same)
        
        //「1つ増加する」パターン        
        dp[n+1][k+1] += dp[n][k] * (n - k - same)
    }
}

• 「変わらない」パターンでは同じ数が続いた分だけ加算し、「1つ増加する」パターンでは同じ数が続いた分だけ減算している。
• 理由は分かるだろうか？.....これは、トンボ式投げっぱなしジャーマン¹は厳しいと思うので、解説するね。
• 例えば挿入先配列Bの状態が[3 5 2 9 4]とき、同じ数としては9が挿入されます。Aはソート済みだからね。
  • 「変わらない」パターンを考えると、置ける場所は、「<」のペアの間(k個)とBの先頭(1個)と9の右隣(1個)があります。ほら、k+1+1(same)になったよ！
  • 「1つ増加する」パターンを考えると、置ける場所は、「≧」のペアの間(n-k-1個)とBの末尾(1個)なはずだけど、「9≧4」の間に9を差し込んでも「<」のペアは増えないから、9の左隣(-1個)です。ほら、n-k-1(same)になったよ！
• 次に、同じ数が既に2つ有る場合として、2つが離れている場合①[3 9 2 9 4]とくっついている場合②[3 2 9 9 4]を考えてみよう。
  • 「変わらない」パターンを考えると、①の場合は9の右隣(2個)がおまけで付いてます。②の場合も9の右隣(2個)だけど...あれ？②の場合9の右隣は2個有るけど、出来上がりの見た目は9が3つ有るだけだから、違いがなくない？と思ったけど、元々の厳密な問題の条件は、最初の並び順を配列を[1 2 3 4 5]としたとき、条件を満たす並び順配列、(例[5 2 1 4 3])はいくつあるか？という問題だから、見た目は9が3つ並んでいるだけで違いがなくてもカウント可能。よって、この場合も結果はk+1+2(same)になるね！
  • 「1つ増加する」パターンは...以下略¹

定番の挿入dp問題の解

• 前項で、dpの書き方まで書いたから、後は流れ作業です。

extension Array { // 2次元配列の省略記法用のエクステンション
    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]
var A = readLine()!.split(separator:" ").map{Int($0)!}
let mod = 998244353

A.sort()

var dp = [[Int]](N+1,N+1,0) // 省略記法
dp[0][0] = 1
var same = 0
for n in 0..<N {
    if n > 0 && A[n-1] == A[n] {same += 1} else {same = 0} 
    for k in 0...K {
        dp[n+1][k] += dp[n][k] * (k + 1 + same) % mod
        dp[n+1][k+1] += dp[n][k] * (n - k - same) % mod
    }
}

print(dp[N][K] % mod)

• コレで提出すると、400ms程度でACです。

最後に

• 作ってみると簡単だけど、AtCoderの公式解説とかちょっと酷いし、公式以外の人の解説も省略が激しくて理解に苦しんだよ。
• 同じように挫けかけた人は、僕の解説で理解してくれると嬉しい。でもまぁ、マイナーすぎるswiftだから、そもそもswiftが分からんよ、と言われちゃうかもだけど。

1. ・・・その先は言う必要ないですよね。 自分で考えてみてください。 ↩ ↩²