LeetCode 337. House Robber III
問題
泥棒が複数の家からお金を盗みます。おかしな街なので二分木の構造をしています。各家はノードで表すことができ、ノードの値がその家から盗めるお金を表しています。泥棒はこの街からなるべく多くのお金を盗もうとしています。ただし、エッジで直接繋がっている２つの家に泥棒に入ると通報されてしまいます。二分木のrootが与えられた時に、泥棒が通報を受けずに盗める金額の最大値を計算してください。
ステップ１：単純に考える
一見するとこの問題は「optimal substructure」の問題に見えます。つまり左と右の木の最大の合計を計算します。
この方針でやってみましょう。まずはルートの左と右の木それぞれについて最大値を求めます。ポイントはsubproblemをいかにつくるかということです。rob(node)という関数がそのノードをルートとする木の最大値を返すとすればrob(root)のsubproblemはrob(root.left)とrob(root.right)です。この２つが分かれば解けます。
明らかに再帰を使うようです。再帰では次の２つのことに注意すべきです。
終了条件：どんな時に計算をせずにrob(node)の答えを見つけられるでしょうか？それはもちろんnodeが空の時です。
再帰条件：どうやってrob(node.left)とrob(node.right)からrob(node)の答えを見つけるかという問題です。nodeに着目するとnodeをとるかとらないかの２パターンがあることがわかります。nodeを取らない場合はnode.left, node.rightのsubproblemになります。nodeをとる場合、nodeに直接連結するnode.left, node.rightは取れないので、自動的にnode.left.left, node.left.right, node.right.left, node.right.rightのsubproblemになります。
実際のコードを見てみましょう。
public int rob(TreeNode root) {
if (root == null) return 0;
int val = 0;
if (root.left != null) {
val += rob(root.left.left) + rob(root.left.right);
}
if (root.right != null) {
val += rob(root.right.left) + rob(root.right.right);
}
return Math.max(val + root.val, rob(root.left) + rob(root.right));
}
しかしこの方法は計算量が爆発して（1186ms）とうていアクセプトされません。
ステップ２：もう一歩踏み込む
ステップ１ではoptimal substructureを考えてみました。しかしよく考えてみるとrob(node)において２パターンに分かれる時、nodeをとるケースでnode.left.left, node.left.right, node.right.left, node.right.rightの４つのケースを探索し、しかもその４ケースにおいてさらに４ケースを考える必要があります。葉の方では何度も同じ計算が行われている（overlapping of subproblems）*ことになります。これが計算量爆発の原因です。そこでDPを使います。DPの簡単な実装としてkeyにnode、valueにrob(node)の答えを書いたハッシュを使います。
＊このようなoptimal substructure + overlapping of subproblemではDPを使います。
public int rob(TreeNode root) {
return robSub(root, new HashMap<>());
}
private int robSub(TreeNode root, Map<TreeNode, Integer> map) {
if (root == null) return 0;
if (map.containsKey(root)) return map.get(root);
int val = 0;
if (root.left != null) {
val += robSub(root.left.left, map) + robSub(root.left.right, map);
}
if (root.right != null) {
val += robSub(root.right.left, map) + robSub(root.right.right, map);
}
val = Math.max(val + root.val, robSub(root.left, map) + robSub(root.right, map));
map.put(root, val);
return val;
}
O(n)のメモリ量（nはノードの数）と引き換えにステップ１から大幅に早くなりました（1186ms -> 9ms）。
ステップ３：一歩引いてみる
ステップ１ではrootからなる木の最適解を見つけるrob(root)を考え、その結果としてDPに行きつきました。
一歩引いて、なぜoverlapping subproblemに当たってしまったのか、その原因を考えてみます。実はその原因はrob(node)を「nodeからなる木の最適解を見つける」と定義してしまったところにあります。上述した通りこの定義ではnodeをとる場合と取らない場合を場合わけする必要があります。この際、rob(node)は２つのケースを自動判別できません。そのため「場合わけに関する情報」が深さに従って失われてしまいます。
情報が失われるとは具体的にどういうことでしょうか。rob(node)を次のように定義してみます。
- rob(node)は２つの要素を持つ配列を返す
- 1つめの要素は「nodeを取らなかった場合の最大値」で２つめの要素は「nodeを取った場合の最大値」
１つ目の要素は（nodeを取らなくていいので）rob(node.left)（が返す配列の２つの要素のうち）の最大値とrob(node.right)の最大値の合計です。２つめの要素はどうなるでしょうか。２つめの要素は「nodeを取った場合」なのでnode.left, node.rightのどちらもとることができません。従って２つめの要素はrob(node.left)の１つめの要素とrob(node.right)の１つ目の要素の合計にnode.valを足した値です。
何をしているかというとrob(node)が返す値に２つのパターン（nodeをとるケースと取らないケース）に関する情報を入れることでsubproblemの重なりを回避し、最終的に貪欲法に落とし込んでいます。
解答：
public int rob(TreeNode root) {
int[] res = robSub(root);
return Math.max(res[0], res[1]);
}
private int[] robSub(TreeNode root) {
if (root == null) return new int[2];
int[] left = robSub(root.left);
int[] right = robSub(root.right);
int[] res = new int[2];
res[0] = Math.max(left[0], left[1]) + Math.max(right[0], right[1]);
res[1] = root.val + left[0] + right[0];
return res;
}
Discussion