🐾

LeetCode 337. House Robber III

2022/06/06に公開

問題

泥棒が複数の家からお金を盗みます。おかしな街なので二分木の構造をしています。各家はノードで表すことができ、ノードの値がその家から盗めるお金を表しています。泥棒はこの街からなるべく多くのお金を盗もうとしています。ただし、エッジで直接繋がっている2つの家に泥棒に入ると通報されてしまいます。二分木のrootが与えられた時に、泥棒が通報を受けずに盗める金額の最大値を計算してください。

ステップ1:単純に考える

一見するとこの問題は「optimal substructure」の問題に見えます。つまり左と右の木の最大の合計を計算します。

この方針でやってみましょう。まずはルートの左と右の木それぞれについて最大値を求めます。ポイントはsubproblemをいかにつくるかということです。rob(node)という関数がそのノードをルートとする木の最大値を返すとすればrob(root)のsubproblemはrob(root.left)とrob(root.right)です。この2つが分かれば解けます。

明らかに再帰を使うようです。再帰では次の2つのことに注意すべきです。

終了条件:どんな時に計算をせずにrob(node)の答えを見つけられるでしょうか?それはもちろんnodeが空の時です。

再帰条件:どうやってrob(node.left)とrob(node.right)からrob(node)の答えを見つけるかという問題です。nodeに着目するとnodeをとるかとらないかの2パターンがあることがわかります。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)とうていアクセプトされません。

ステップ2:もう一歩踏み込む

ステップ1ではoptimal substructureを考えてみました。しかしよく考えてみるとrob(node)において2パターンに分かれる時、nodeをとるケースでnode.left.left, node.left.right, node.right.left, node.right.rightの4つのケースを探索し、しかもその4ケースにおいてさらに4ケースを考える必要があります。葉の方では何度も同じ計算が行われている(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はノードの数)と引き換えにステップ1から大幅に早くなりました(1186ms -> 9ms)。

ステップ3:一歩引いてみる

ステップ1ではrootからなる木の最適解を見つけるrob(root)を考え、その結果としてDPに行きつきました。

一歩引いて、なぜoverlapping subproblemに当たってしまったのか、その原因を考えてみます。実はその原因はrob(node)を「nodeからなる木の最適解を見つける」と定義してしまったところにあります。上述した通りこの定義ではnodeをとる場合と取らない場合を場合わけする必要があります。この際、rob(node)は2つのケースを自動判別できません。そのため「場合わけに関する情報」が深さに従って失われてしまいます。

情報が失われるとは具体的にどういうことでしょうか。rob(node)を次のように定義してみます。

  • rob(node)は2つの要素を持つ配列を返す
  • 1つめの要素は「nodeを取らなかった場合の最大値」で2つめの要素は「nodeを取った場合の最大値」

1つ目の要素は(nodeを取らなくていいので)rob(node.left)(が返す配列の2つの要素のうち)の最大値とrob(node.right)の最大値の合計です。2つめの要素はどうなるでしょうか。2つめの要素は「nodeを取った場合」なのでnode.left, node.rightのどちらもとることができません。従って2つめの要素はrob(node.left)の1つめの要素とrob(node.right)の1つ目の要素の合計にnode.valを足した値です。

何をしているかというとrob(node)が返す値に2つのパターン(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

ログインするとコメントできます