Problem

The thief has found himself a new place for his thievery again. There is only one entrance to this area, called root.

Besides the root, each house has one and only one parent house. After a tour, the smart thief realized that all houses in this place form a binary tree. It will automatically contact the police if two directly-linked houses were broken into on the same night.

Given the root of the binary tree, return the maximum amount of money the thief can rob without alerting the police.

https://leetcode.com/problems/house-robber-iii/

Example 1:

case1

Input: root = [3,2,3,null,3,null,1]
Output: 7
Explanation: Maximum amount of money the thief can rob = 3 + 3 + 1 = 7.

Example 2:

case2

Input: root = [3,4,5,1,3,null,1]
Output: 9
Explanation: Maximum amount of money the thief can rob = 4 + 5 = 9.

Constraints:

  • The number of nodes in the tree is in the range [1, 10⁴].
  • 0 <= Node.val <= 10⁴

Test Cases

1
2
3
4
5
6
7
8
# Definition for a binary tree node.
# class TreeNode:
#     def __init__(self, val=0, left=None, right=None):
#         self.val = val
#         self.left = left
#         self.right = right
class Solution:
    def rob(self, root: Optional[TreeNode]) -> int:
solution_test.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
import pytest

import os
import sys
sys.path.append(os.path.dirname(os.path.dirname(__file__)))
from _utils.binary_tree import build_tree
from solution import Solution
from solution2 import Solution as Solution2
from solution3 import Solution as Solution3

null = None


@pytest.mark.parametrize('root, expected', [
    ([3,2,3,null,3,null,1], 7),
    ([3,4,5,1,3,null,1], 9),
])
@pytest.mark.parametrize('sol', [Solution(), Solution2(), Solution3()])
class Test:
    def test_solution(self, sol, root, expected):
        assert sol.rob(build_tree(root)) == expected

Thoughts

系列题：

本题是从数组改成了二叉树。

稍微扩展一下 198. House Robber 中定义的 t[i]，改为 t(u)（其中 u 是二叉树中的一个节点），表示从以 u 为根节点的子树中，能抢到的最多的金额（注意房间 u 可能抢也可能不抢）。初值 t(null) = 0，题目所求结果为 t(root)。

对于节点 u，可以选择抢，得到金额 u.val，那就不能抢 u.left 和 u.right，但是这两个子节点的再下一层子节点就都可以抢了，总额为 u.val + t(u.left.left) + t(u.left.right) + t(u.right.left) + t(u.right.right)。

也可以选择不抢，那么 u.left 和 u.right 就都可以抢，总额为 t(u.left) + t(u.right)。

所以：

t(u)=max⁡{u.val+t(u.l.l)+t(u.l.r)+t(u.r.l)+t(u.r.r)t(u.l)+t(u.r)t(u)=\max\begin{cases} u.val+t(u.l.l)+t(u.l.r)+t(u.r.l)+t(u.r.r) \\ t(u.l)+t(u.r) \end{cases}

按后序（post-order，LRN）遍历二叉树，计算每个节点的 t。非递归的后序遍历可以用栈辅助，参考 124. Binary Tree Maximum Path Sum。

时间复杂度 O(n)，空间复杂度 O(log n)（遍历所用的栈空间）（t 值可以直接存在节点的 val 里）。

Code

solution.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
from typing import Optional


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    def rob(self, root: Optional[TreeNode]) -> int:
        if not root: return 0

        max2 = lambda a, b: a if a >= b else b
        val = lambda u: u.val if u else 0
        stack = []
        prev = None
        node = root
        while node or stack:
            if node:
                stack.append(node)
                node = node.left
            elif stack[-1].right != prev:
                node = stack[-1].right
                prev = None
            else:
                prev = stack.pop() # prev is the LRN visiting node, its left and right child both done.
                if prev.left:
                    prev.val += val(prev.left.left) + val(prev.left.right)
                if prev.right:
                    prev.val += val(prev.right.left) + val(prev.right.right)
                prev.val = max2(prev.val, val(prev.left) + val(prev.right))

        return root.val

Another DP

如果嫌访问第二层子节点会有太多是否为空的判定，也可以改造一下状态变量的定义。其实类似于 198. House Robber 中定义的 第二种 DP，明确抢还是不抢当前房间。

定义 ty(u) 表示抢房间 u 的最大总额，tn(u) 表示不抢房间 u 的最大总额。

如果选择抢房间 u，那么 u.left 和 u.right 就都不能抢，ty(u) = u.val + tn(u.left) + tn(u.right)。

如果选择不抢房间 u，那么 u.left 和 u.right 就都可抢可不抢，哪个收益大选哪个，于是 tn(u) = max{ty(u.left), tn(u.left)} + max{ty(u.right), tn(u.right)}。

所以：

{ty(u)=u.val+tn(u.l)+tn(u.r)tn(u)=max⁡{ty(u.l),tn(u.l)}+max⁡{ty(u.r),tn(u.r)}\begin{cases} ty(u)=u.val+tn(u.l)+tn(u.r) \\ tn(u)=\max\{ty(u.l),tn(u.l)\}+\max\{ty(u.r),tn(u.r)\} \end{cases}

最终结果取 max{ty(root), tn(root)}。

solution2.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
from typing import Optional


# Definition for a binary tree node.
class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right


class Solution:
    def rob(self, root: Optional[TreeNode]) -> int:
        if not root: return 0

        max2 = lambda a, b: a if a >= b else b
        dp = lambda u: u.val if u else (0, 0)
        stack = []
        prev = None
        node = root
        while node or stack:
            if node:
                stack.append(node)
                node = node.left
            elif stack[-1].right != prev:
                node = stack[-1].right
                prev = None
            else:
                prev = stack.pop() # prev is the LRN visiting node, its left and right child both done.
                prev.val = (
                    prev.val + dp(prev.left)[1] + dp(prev.right)[1],
                    max2(*dp(prev.left)) + max2(*dp(prev.right)),
                )

        return max2(*root.val)

PS：上边是记录 ty(u) 和 tn(u) 这两个值，也可以考虑记录 t(u) 和 tn(u)，其中 t(u) = max{ty(u), tn(u)}，代码见 solution3.py。

medium
参考资料
许可协议

本文除了题目描述（Problem）部分之外，其他内容由 Calf 原创，采用 CC BY-NC-SA 4.0 许可协议，转载请注明出处。

分享文章

快来参与讨论吧~