はじめに
木dpの考え方について
- 木において1つの頂点を根として設定することがある。このとき、根以外の頂点は、根に向かう経路をただ1本持つ。このとき、その頂点を「子」と考えると、根に向かう経路上に隣接する頂点を「親」と呼ぶ。
- ある頂点iから、頂点iの親に向かう辺を削除し、頂点iを根とする木を考えると、その木は頂点iを根とする部分木と呼べる。
- ここまで書けば、想像つくと思うけど、頂点iを根と見なした部分木における解をdp[i][何かの状態]として、元々の根が頂点0の時、最終的にdp[0][･･･]が解となる。
- このような考え方だから、木dpはメモ化再帰で解くのが定番っぽいね。
問題例
- サイトで出会った木dpはこれです。
- 問題の内容はシンプルです。
- 各頂点を「白」または「黒」に塗ることが出来ます。
- 隣り合う頂点同士を「黒」に塗ることは出来ません。
- 頂点を塗るパターンはいくつあるか。(mod 1000000007)
- 今回も生成AIに教えを請うことにしました。標準入力部分はうまく伝えられなかったから、自作しました。
- メモ化dp配列[node][0/1]は、頂点nodeを根とした部分木について、頂点nodeを白(0)または黒(1)に塗ったときの塗るパターン数です。
- 再帰関数dfs(node,parent)は、頂点nodeを根と見なす部分木における解をdp[node][0]とdp[node][1]に書き込む関数です。また、parentはnodeから見た親の頂点です。
- 初期値として、dpの全ての値を1としておきます。
// import Foundation -- 不要なのでコメントアウト
/// 標準入力対応は自作
let n = Int(readLine()!)!
var edges:[(Int,Int)] = []
for _ in 0..<n-1 {
let ab = [readLine()!.split(separator:" ").map{Int($0)!}].map{($0[0],$0[1])}[0]
edges.append(ab)
}
let MOD = 1000000007
func solveTreeDP(_ n: Int, _ edges: [(Int, Int)]) -> Int {
var graph = [[Int]](repeating: [], count: n)
for (u, v) in edges {
graph[u - 1].append(v - 1)
graph[v - 1].append(u - 1)
}
var dp = [[Int]](repeating: [1, 1], count: n)
func dfs(_ node: Int, _ parent: Int) {
for neighbor in graph[node] {
if neighbor == parent {
continue
}
dfs(neighbor, node)
dp[node][0] = (dp[node][0] * (dp[neighbor][0] + dp[neighbor][1])) % MOD
dp[node][1] = (dp[node][1] * dp[neighbor][0]) % MOD
}
}
dfs(0, -1)
return (dp[0][0] + dp[0][1]) % MOD
}
// let n = 5 // 頂点数の例 -- 生成AIの作った入力例
// let edges = [(1, 2), (1, 3), (3, 4), (3, 5)] // 辺の例 -- 生成AIの作った入力例
print(solveTreeDP(n, edges))
- 再帰呼び出し後のメモ化配列の計算は2種類あります。
- dp[node][0] = (dp[node][0] * (dp[neighbor][0] + dp[neighbor][1]))
- 頂点nodeの色が白いときは、子neighborの色は白と黒の両方が可能です。
- dp[node][1] = (dp[node][1] * dp[neighbor][0])
- 頂点nodeの色が黒いときは、子neighborの色は白のみです。
- dp[node][0] = (dp[node][0] * (dp[neighbor][0] + dp[neighbor][1]))
- dfsの大元の呼び出しは、根(node:0)に対するものですが、根に親はないので、親を-1(頂点番号に存在しないインデックス)としています。
- 解答は、根の色が黒い場合と白い場合の合計となっています。
- これで提出すると、150ms程度でACでした。
さいごに
- 前回の投稿に続いて、生成AIを頼ってしまいました。生成AIは解答だけでなく簡単な解説までしてくれてます。
- 全くの初心者だと、簡単な解説だけじゃ意味不明だろうけど、今回で言えば、グラフ理論における木やdpの初歩を知っていれば、生成AIの解説がなくても、コードだけ見れば何がやりたいか分かる。
- これまでは、初めての解法に出会ったとき、誰かの書いた分かりやすい解法を探すのが大変だったけど、これからは生成AIに解き方を教えて貰えばいいのかな？
- 生成AIが最初に大打撃を与える職業は学校の先生な気がしてきた。