LeetCode 337: House Robber III

Problem Restatement

We are given the root of a binary tree.

Each node represents a house, and the value of the node is the amount of money in that house.

A thief wants to rob houses, but there is one rule:

If two directly connected houses are robbed, the police are alerted.

In tree terms, we cannot rob both a parent node and its child node.

Return the maximum amount of money that can be robbed.

The problem constraints say the tree has between 1 and 10^4 nodes, and each node value is between 0 and 10^4.

Input and Output

Function shape:

def rob(root: Optional[TreeNode]) -> int:
    ...

Examples

Example 1:

Input: root = [3,2,3,null,3,null,1]
Output: 7

The tree is:

      3
     / \
    2   3
     \   \
      3   1

The best choice is:

3 + 3 + 1 = 7

We rob the root and the two grandchildren.

Example 2:

Input: root = [3,4,5,1,3,null,1]
Output: 9

The tree is:

      3
     / \
    4   5
   / \   \
  1   3   1

The best choice is:

4 + 5 = 9

We skip the root and rob its children.

First Thought: Recursion With Choices

At each node, we have two choices:

A naive recursive solution may repeatedly recompute the same subtrees.

For example, when deciding whether to rob the root, we may compute values for children and grandchildren many times.

We need each subtree to give its parent enough information in one pass.

Key Insight

For every node, compute two values:

If we rob the current node, we cannot rob its direct children.

So:

rob_node = node.val + left.skip + right.skip

If we skip the current node, each child can choose its better option.

So:

skip_node = max(left.rob, left.skip) + max(right.rob, right.skip)

This is dynamic programming on a tree.

Algorithm

Use postorder DFS.

For each node:

1. Recursively compute (left_rob, left_skip) for the left child.
2. Recursively compute (right_rob, right_skip) for the right child.
3. Compute the value if robbing the current node.
4. Compute the value if skipping the current node.
5. Return both values to the parent.

For a null node:

rob = 0
skip = 0

At the end, the answer is:

max(root_rob, root_skip)

Walkthrough

Use:

      3
     / \
    2   3
     \   \
      3   1

Leaf node 3:

rob = 3
skip = 0

Leaf node 1:

rob = 1
skip = 0

Node 2 has no left child and right child 3.

If we rob 2:

rob = 2 + 0 + 0 = 2

If we skip 2:

skip = max(0, 0) + max(3, 0) = 3

So node 2 returns:

(2, 3)

Node 3 on the right has right child 1.

If we rob it:

rob = 3 + 0 + 0 = 3

If we skip it:

skip = max(0, 0) + max(1, 0) = 1

So it returns:

(3, 1)

Now process the root 3.

If we rob the root:

rob = 3 + left.skip + right.skip
    = 3 + 3 + 1
    = 7

If we skip the root:

skip = max(2, 3) + max(3, 1)
     = 3 + 3
     = 6

The answer is:

max(7, 6) = 7

Correctness

For each node, there are only two possible states relevant to its parent: robbed or skipped.

If the current node is robbed, then its children cannot be robbed. The best value in that case is the current node value plus the skip values of both children.

If the current node is skipped, then each child is independent. For each child, we may choose whichever state gives more money.

The DFS computes these two values after computing the same two values for the children. This postorder traversal ensures every subtree is solved before its parent uses it.

The returned pair for a node therefore gives the optimal robbery amount for that subtree under both possible states of the node.

The root has no parent, so either state is allowed. Taking the maximum of the two root states gives the maximum amount of money that can be robbed without robbing any directly connected pair.

Complexity

In the worst case, a skewed tree has height n, so recursion space can be O(n).

Implementation

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:
        def dfs(node: Optional[TreeNode]) -> tuple[int, int]:
            if node is None:
                return 0, 0

            left_rob, left_skip = dfs(node.left)
            right_rob, right_skip = dfs(node.right)

            rob_node = node.val + left_skip + right_skip

            skip_node = (
                max(left_rob, left_skip)
                + max(right_rob, right_skip)
            )

            return rob_node, skip_node

        root_rob, root_skip = dfs(root)

        return max(root_rob, root_skip)

Code Explanation

The DFS returns two values:

return rob_node, skip_node

For a missing child:

if node is None:
    return 0, 0

A null subtree contributes no money.

First compute the child states:

left_rob, left_skip = dfs(node.left)
right_rob, right_skip = dfs(node.right)

If we rob the current node, we must skip both children:

rob_node = node.val + left_skip + right_skip

If we skip the current node, each child can choose its better state:

skip_node = max(left_rob, left_skip) + max(right_rob, right_skip)

The final answer is the better root state:

return max(root_rob, root_skip)

Testing

class TreeNode:
    def __init__(self, val=0, left=None, right=None):
        self.val = val
        self.left = left
        self.right = right

def run_tests():
    s = Solution()

    root = TreeNode(
        3,
        TreeNode(2, None, TreeNode(3)),
        TreeNode(3, None, TreeNode(1)),
    )
    assert s.rob(root) == 7

    root = TreeNode(
        3,
        TreeNode(4, TreeNode(1), TreeNode(3)),
        TreeNode(5, None, TreeNode(1)),
    )
    assert s.rob(root) == 9

    root = TreeNode(5)
    assert s.rob(root) == 5

    root = TreeNode(
        2,
        TreeNode(1),
        TreeNode(3),
    )
    assert s.rob(root) == 4

    root = TreeNode(
        4,
        TreeNode(1, TreeNode(2), TreeNode(3)),
        TreeNode(5),
    )
    assert s.rob(root) == 14

    print("all tests passed")

run_tests()

Test meaning: