教育的な「メモ化+再帰関数」のdp問題【競技プログラミング】

はじめに

• 最近、AtCoderのDPまとめコンテストに取り組んでいるけど、遷移型のDPとメモ化再帰（メモ化＋再帰関数）のDPが混ざっているんだよね。まぁ、動的計画法が両方を含む概念だとは分かっているけど、頭がこんがらがるよね。
• そもそも、最初に動的計画法に出会ったのは、メモ化再帰の方だったけど、初心者の頃は、遷移型の方が分かり易くて良いよね。というか、そもそも動的計画法という同じ名称で括ってしまうこと自体が間違っていると思う。
• で、今回は、メモ化再帰として、簡単すぎず、難しすぎない問題を「DPまとめコンテスト」の中から紹介しようと思う。

教育的なメモ化再帰の問題

• 紹介するのは、問題Ｌ:Dequeです。
• 内容は、つぎのとおり、
  • 太郎･次郎の二者によるゲームで、共に点数を最大化するための行動を行うものとする。
  • 長さNの整数配列の右端･左端のどちらか好きな方から交互に数字(これが点数になる)を抜き取り、配列の要素がなくなるまで繰り返す。
  • 太郎が先行の時、「太郎の点数合計X-次郎の点数合計Y」はいくらとなるか？
• 入力例

4 // 要素数は4個
10 80 90 30 // 4個の整数列

• 上記の場合、太郎30 ⇒ 次郎90 ⇒ 太郎80 ⇒ 次郎10の順でゲットするから、解は10(= 太郎110 - 次郎100)。

再帰関数について

• メモ化＋再帰関数のキモは、当然、『いい感じの再帰関数』を作ること。
• 例えばこんな感じ。ちなみに関数名recはrecursive(再帰的)の略。配列Aは与えられた長さNの整数列(例えば、長さ4の[10 80 90 30])とする。

func rec(_ l:Int ,_ r:Int) ->Int {
    // 終端条件（ベースケース）
    if l == r {return 0}

    // 再帰ステップ
    let ans = max(A[l] - rec(l+1,r) , A[r-1] - rec(l,r-1)) // 再帰呼び出し
    return ans
}

• 上記の再帰関数は、配列Aの部分配列(l..<r)の時の解なので、rec(0,N)が求める解となります。
• 再帰関数が「終端条件」と「再帰ステップ」で成り立つのは理解してるかな？理解していない人は、この投稿のフィボナッチ関数の解き方を見てくれれば分かるかな？
• 上記の再帰ステップは、分かり易いよね。左端をゲットしたときと、右端をゲットしたときの大きい方を解答しているだけ。
  • 左端をゲットしたときの点数がA[l]-rec(l+1,r)となっているけど、A[l]は太郎がゲットした左端の点数であり、-rec(l+1,r)が相手方(次郎)の最適解の控除となる。
• 再帰関数って、出来上がりを見ると凄く単純だけど、これが頭に思い浮かぶかどうかが、運命の分かれ目。
• 多くの問題をこなして、頭の中に再帰関数のパターンが染み込むようにした方が効率が良いので、修行期間中は、思い浮かばなかったらすぐに解法を見た方が良い。
  • 答えをすぐに見ちゃうと深く考える力がつかないかも知れないけど、ABCコンテストに定期的に参加していれば大丈夫でしょ。
• 上記の再帰関数を使えば、計算量の少ない問題は解ける。

let N = Int(readLine()!)!
let A = readLine()!.split(separator:" ").map{Int($0)!}

print(rec(0,N))

func rec(_ l:Int ,_ r:Int) ->Int {
    if l == r {return 0}
    let ans = max(A[l] - rec(l+1,r) , A[r-1] - rec(l,r-1))
    return ans
}

• 上記を提出するとTLEになるけど、７問(7/19)は解けたよ。

メモ化について

• 再帰関数で問題は解くことが出来たけど、TLEになってしまったのは、再帰呼び出しで何度も同じ計算をしちゃってたから。
• メモ化の意味は、過去の投稿でも触れているから、それを読んで欲しい。
• メモを格納する配列名を再帰計画法らしくdp[l][r]（rec(l,r)の答えを格納）とすると、再帰関数は下記の通りとなる。

func rec(_ l:Int ,_ r:Int) ->Int {
    if dp[l][r] != INF {return dp[l][r]} // ①メモを利用
    if l == r {return 0}
    
    let ans = max(A[l] - rec(l+1,r) , A[r-1] - rec(l,r-1))
    dp[l][r] = ans // ②メモする
    return ans
}

• ②でメモして、①で利用する。上記のコードだけ見ると、順序が逆じゃん！と思うけど、再帰ステップで利用されるから意味があるよ。
  • ピンとこない人は、過去の投稿でのフィボナッチ数列の再帰的解法を自分でコードを書いて、スルメのように噛みしめて見た方が良いよ。
  • 過去の投稿では、メモ化の配列を２つ(mems:メモ化されているか？、fibs:メモの内容)に分けているけど、メモ化の配列を1つ(dp:初期値-1とし、値が-1の時はメモ化前、値が0以上ならメモの内容)にするように書き換えてみたら良い練習になるよ。

全体解答

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 = Int(readLine()!)!
let A = readLine()!.split(separator:" ").map{Int($0)!}

let INF = -Int.max / 2 // -Int.maxだと、違う問題だとオーバーフローでエラーしちゃう可能性があるから、INFを使うときはいつも半分にしてる。この問題だと、半分にする必要ないけどね。要は、偶然一致しないような数字にしておけば良いだけ。swiftだから、オプショナル値で処理しても良いけど、コードが冗長になるのよね。
var dp = [[Int]](N+1,N+1,INF)

print(rec(0,N))

func rec(_ l:Int ,_ r:Int) ->Int {
    if dp[l][r] != INF {return dp[l][r]} // ①メモを利用
    if l == r {return 0}
    
    let ans = max(A[l] - rec(l+1,r) , A[r-1] - rec(l,r-1))
    dp[l][r] = ans // ②メモする
    return ans
}

• 投稿したら、無事にACになりました。

おわりに

• 再帰関数が短く書けて、まさに教育的な問題だったね。
• そしておそらく、メモ化再帰以外のアルゴリズムでこの問題を解くのは難しいだろうね。少なくとも僕は思いつかない。
• この調子で「DPまとめコンテスト」を全クリして、レベルアップを図ろうと思う。

おまけ

• 上記回答では、太郎の点数Xと次郎の点数Yの差額を求めたけど、X,Y自体を求めるにはどうすれば良いだろうか？
• メモ化配列dpに、配列Aの左右どちらからゲットしたかの履歴を残せばよい。
• 具体的には、以下のような改造を行う。

// メモ化配列の改造
var dp = [[(Int,Bool?)]](N+1,N+1,(INF,nil)) // タプルのセカンドをBoolとし、配列の左からゲットした場合はtrue、右からゲットしたときは、falseとする。

// max関数の改造
func maxLR(_ l:Int,_ r:Int)->(Int,Bool) {
    (max(l,r),l >= r) // 2引数のうち、左側の方が大きいとき、タプルのセカンドでtrueを返す。
}

// 再帰関数の改造
func rec(_ l:Int ,_ r:Int) ->Int {
    if dp[l][r].0 != INF {return dp[l][r].0}
    if l == r {return 0}
    
    let (ans,lr) = maxLR(A[l] - rec(l+1,r) , A[r-1] - rec(l,r-1))
    dp[l][r].0 = ans
    dp[l][r].1 = lr // 左右どちらからゲットしたかの履歴を残す
     
    return ans
}

• 上記改造後、履歴をゲットし、X,Yを求めるコード。

print(rec(0,N)) // 元々の回答（この再帰関数を実行しないと、メモ化がされないので必須）

// メモ化後の左右履歴(lr)の活用
var (l,r) = (0,N)
var (x,y) = (0,0)
var taro = true
while l != r {
    let lr = dp[l][r].1! // 左右履歴の取得
    let val = lr ? A[l] : A[r-1]
    if taro {x += val} else {y += val}
    if lr {l += 1} else {r -= 1} 
    taro = !taro
 }

print("X =",x)
print("Y =",y)

• 入力例

6
4 2 9 7 1 5

• 出力例

2
X = 15
Y = 13

おまけ -- その２

• 全然違う問題だけど、教育的な問題を見つけたから、投稿するね。
• 元の問題
  • N個のスライムが並んでいて、スライムiの重さはA[i]で与えられる。
  • 隣同士のスライムは合体させることが出来て、コストは２つのスライムの重さの合計。
  • スライムを全部合体させるときの最低コストを答えよ。
• メモ化再帰で解答しやすいように、下記の問題とみなす。
• 見なし問題
  • 全て合体した後の重さΣA[i]のスライムを、最終的にN個のスライムに切り分ける。
  • スライムを2つに切り分けるときのコストは、元のスライムの重さとする。
• 見なし問題を解くのに最適な再帰関数
  • rec(l,r) -- スライムl...rが合体したスライムを全て切り分けるコスト。
• 解答は以下のとおり

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 = Int(readLine()!)!
let A = readLine()!.split(separator:" ").map{Int($0)!}

// Aの累積和配列を作成
var sumA = [0]
var sum = 0
for a in A {
    sum += a
    sumA.append(sum)
}

// スライムの合体範囲がl...rのときの分解コストのメモ化配列dp[l][r]
var dp = [[Int]](N,N,-1)

print(rec(0,N-1))

// スライムの合体範囲がl...rのときの分解コストの再帰関数rec(l,r)
func rec(_ l:Int,_ r:Int) -> Int{

    // 終端条件
    if l == r {return 0}
    if dp[l][r] != -1 {return dp[l][r]}

    // 再帰ステップ
    var div = sumA[r+1] - sumA[l] // 2分割のコスト
    var cost = Int.max / 2 // ２分割したあとの２つのスライムを最後まで分割するコスト
    for m in l..<r { // 最もコストが安くなる切り口を選ぶ
        cost = min(cost,rec(l,m) + rec(m+1,r)) // 再帰呼び出し
    }
    let ans = div + cost
        
    dp[l][r] = ans // メモ化
    return ans 
}

• 上記コードは結構簡単そうに見えるけど、見落としがちな観点は、2分割した後の２つのスライムを最後まで分割するコスト（cost）を算出する下記のコード

    for m in l..<r { // 最もコストが安くなる切り口を選ぶ
        cost = min(cost,rec(l,m) + rec(m+1,r)) // 再帰呼び出し
    }

• 上記for文において、再帰呼び出しは「rec(l,l)＋rec(l+1,r)」と「rec(l,r-1)+rec(r,r)」を含む。
  • 前者は、スライム[l]とスライム[l+1...r]への２分割だけど、スライム[l]はこれ以上分割できないから、スライム[l]の分割コストrec(l,l)は0となる。これは、終端条件のif l == r {return 0}で得られる。
  • 後者のrec(r,r)も終端条件で算出されるね。このような形で終端条件が使われることについて、パッと思い浮かばないと、切口mを選ぶときに、for m in l+1..<r-1等と書いてしまいエラーを生じさせることとなる。まあ、僕のことなんだけど...
• ちなみに、上記を提出すると約60msでACとなったよ。