See all coding puzzles and patterns here.

Problem Definition

A house robber wants to rob as much money as possible from a binary tree of houses. If the security system will alert the police if two linked houses (in the binary tree) are robbed, determine how much money the robber can rob from this tree. Assume that all houses have a zero or positive amount of money.


Node Class

The class representing each element (a house) in the binary tree of houses.

class Node:
    def __init__(self, val: int, left: 'Node' = None, right: 'Node' = None):
        self.val = val
        self.left = left
        self.right = right

Examples

Some test cases to verify our solution.

# For this tree, the robber can get 7 by robbing the root and the two leaves or 5 by robbing the middle layer
#     3
#    / \
#   2   3
#  /   /
# 3   1
three = Node(3)
one = Node(1)
two = Node(2, left=three)
three_2 = Node(3, left=one)
root = Node(3, left=two, right=three_2)

assert rob(root) == 7

#       4
#      /
#     1
#    /
#   2
#  /
# 3
three = Node(3)
two = Node(2, left=three)
one = Node(1, left=two)
root = Node(4, left=one)

assert rob(root) == 7

Solution

This is a graph traversal problem that can be solved with a breadth-first search or a depth-first search. Here, we implement a DFS solution.

For every node, we will compute two values:

  • How much money can be robbed if the current house is robbed.
  • How much money can be robbed if the current house is not robbed.

We will use the information from children nodes to determine how much money can be robbed if the house is robbed or if the house is not robbed.

For the current node, the max sum could be in these cases:

  1. by robbing the current house and then skipping the next house
  2. by skipping the current house and:
    1. skipping the left house but robbing the right house
    2. skipping the right house but robbing the left house
    3. robbing both the left and right house
    4. skipping both the left and right house
def rob(root: Node) -> int:
    def helper(node: Node) -> tuple[int, int]:
        """
        Returns (max by robbing current node, max by not robbing current node)
        """
        if node.left is None:
            left_rob, left_not_rob = 0, 0
        else:
            left_rob, left_not_rob = helper(node.left)
        if node.right is None:
            right_rob, right_not_rob = 0, 0
        else:
            right_rob, right_not_rob = helper(node.right)
        max_rob = node.val + left_not_rob + right_not_rob
        max_skip = max(left_rob, left_not_rob) + max(right_rob, right_not_rob)
        return (max_rob, max_skip)
    return max(helper(root))

Complexity

Let $n = $ the number of nodes (houses) in the binary tree.

Runtime complexity should be $\Theta(n)$ since we are visiting every node once.

Auxiliary space complexity should be $O(n)$ since, in the worst case, the tree is a linked list, resulting in a worst case of $n$ program stack frames.