337. House Robber III
On LeetCode ->Problem¶
Choose houses in a binary tree to maximize the money robbed, without choosing both a parent and its child.
Key trick¶
For each node, compute two results after processing its children: the best total when skipping it and when robbing it.
Trap¶
A node cannot be calculated when it is first encountered: its children's results must already exist. Also, a tree with 10,000 nodes can be too deep for Python recursion.
Why is it interesting?¶
It turns a local restriction between adjacent houses into a two-state dynamic program on a tree.
Python solution¶
class Solution:
def rob(self, root: "TreeNode | None") -> int:
if root is None:
return 0
# dp[node] = (best if skipped, best if robbed).
dp = {None: (0, 0)}
stack = [(root, False)]
while stack:
node, expanded = stack.pop()
if not expanded:
# Revisit the node only after both children are processed.
stack.append((node, True))
if node.right:
stack.append((node.right, False))
if node.left:
stack.append((node.left, False))
continue
left_skip, left_rob = dp[node.left]
right_skip, right_rob = dp[node.right]
skip = max(left_skip, left_rob) + max(right_skip, right_rob)
rob = node.val + left_skip + right_skip
dp[node] = (skip, rob)
return max(dp[root])
Time: \(O(n)\). Space: \(O(n)\).
Comment on my solution¶
Your two-state recurrence and iterative postorder approach are right. Both versions need a continue after pushing an unexpanded node; otherwise they read child results before those results exist. The second version also needs to assign dp[node] before accessing dp[node][0], and its cleanup deletes the left child twice instead of deleting the right child.
# WRONG
# KeyError: TreeNode{val: 2, left: None, right: TreeNode{val: 3, left: None, right: None}}
# ~~~~~^^^^^^^^^^^
# max(dp[0][node.left], dp[1][node.left]) +
# Line 20 in rob (Solution.py)
# ret = Solution().rob(param_1)
# Line 57 in _driver (Solution.py)
# ~~~~~~~^^
# _driver()
# Line 72 in <module> (Solution.py)
class Solution:
def rob(self, root: TreeNode | None) -> int:
if root is None:
return 0
dp = [{None:0}, {None:0}] # [skip, rob]
stack = [(root, False)] # (node, expanded)
while stack:
node, expanded = stack.pop()
if not expanded:
stack.append((node, True))
if node.right:
stack.append((node.right, False))
if node.left:
stack.append((node.left, False))
# skip node
dp[0][node] = (
max(dp[0][node.left], dp[1][node.left]) +
max(dp[0][node.right] , dp[1][node.right])
)
dp[1][node] = node.val + dp[0][node.left] + dp[0][node.right]
if node.left:
del dp[0][node.left]
del dp[1][node.left]
if node.right:
del dp[0][node.right]
del dp[1][node.right]
return max(dp[0][root], dp[1][root])
# WRONG (also)
class Solution:
def rob(self, root: TreeNode | None) -> int:
if root is None:
return 0
dp = {None: [0,0]} # node -> [max if skipped, max if robbed]
stack = [(root, False)] # (node, expanded)
while stack:
node, expanded = stack.pop()
if not expanded:
stack.append((node, True))
if node.right:
stack.append((node.right, False))
if node.left:
stack.append((node.left, False))
# skip node
dp[node][0] = (
max(dp[node.left][0], dp[node.left][1]) +
max(dp[node.right][0] , dp[node.right][1])
)
dp[node][1] = node.val + dp[node.left][0] + dp[node.right][0]
if node.left:
del dp[node.left]
del dp[node.left]
if node.right:
del dp[node.right]
del dp[node.right]
return max(dp[root])
# WORKS - written after AI comments
class Solution:
def rob(self, root: TreeNode | None) -> int:
if root is None:
return 0
dp = {None: [0,0]} # node -> [max if skipped, max if robbed]
stack = [(root, False)] # (node, expanded)
while stack:
node, expanded = stack.pop()
if not expanded:
stack.append((node, True))
if node.right:
stack.append((node.right, False))
if node.left:
stack.append((node.left, False))
continue
# skip node
dp[node] = [0, 0]
dp[node][0] = (
max(dp[node.left][0], dp[node.left][1]) +
max(dp[node.right][0] , dp[node.right][1])
)
dp[node][1] = node.val + dp[node.left][0] + dp[node.right][0]
if node.left:
del dp[node.left]
if node.right:
del dp[node.right]
return max(dp[root])