DSA House Robber III

The problem

Houses stand at the nodes of a binary tree. Each house holds a value. Robbing a house blocks its parent and its children, because neighbouring houses are linked by alarms.

No two robbed houses may be neighbours in the tree. The answer is the largest total that one night of robbing can take. Houses hold non-negative values in this article.

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

tree_a = TreeNode(3,
                  TreeNode(2, None, TreeNode(3)),
                  TreeNode(3, None, TreeNode(1)))
tree_b = TreeNode(3,
                  TreeNode(4, TreeNode(1), TreeNode(3)),
                  TreeNode(5, None, TreeNode(1)))
tree_c = TreeNode(1, TreeNode(100), TreeNode(2, None, TreeNode(100)))
for name, expected in [("Tree A", 7), ("Tree B", 9), ("Tree C", 200), ("single house 5", 5)]:
    print(name, "-> best haul:", expected)

Tree C is the teaching case. Its root is worth almost nothing, and its two rich leaves sit in different branches. The best night robs both leaves and skips everything between them.

Intuition: rob level by level

A first idea groups houses by depth. Whole levels are robbed or skipped as blocks, alternating down the tree. Two totals compete, the even levels and the odd levels.

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

tree = TreeNode(1, TreeNode(100), TreeNode(2, None, TreeNode(100)))
levels = []
level = [tree]
while level:
    levels.append(sum(node.value for node in level))
    nxt = []
    for node in level:
        if node.left is not None:
            nxt.append(node.left)
        if node.right is not None:
            nxt.append(node.right)
    level = nxt
even = sum(levels[0::2])
odd = sum(levels[1::2])
print("level sums:", levels)
print("even levels total:", even, " odd levels total:", odd)
print("the level method answers:", max(even, odd))
print("the true best haul is 200")

The level method answers 102, and the true best is 200. The winning night mixes levels inside different branches. It robs the left leaf at level 1 and the right leaf at level 2, a mix no whole-level choice can express.

Brute force: try every subset of houses

The brute force lists every subset of the houses. A subset is legal when no house in it has its parent also in it. The best legal total is the answer.

A counter records the subsets examined. The count doubles with each added house, which is the cost this method cannot escape.

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

def brute(root):
    nodes = []

    def collect(node, parent):
        if node is None:
            return
        index = len(nodes)
        nodes.append((node.value, parent))
        collect(node.left, index)
        collect(node.right, index)

    collect(root, -1)
    best = 0
    examined = 0
    for mask in range(2 ** len(nodes)):
        examined += 1
        total = 0
        legal = True
        for i, (value, parent) in enumerate(nodes):
            if mask >> i & 1:
                if parent >= 0 and mask >> parent & 1:
                    legal = False
                    break
                total += value
        if legal and total > best:
            best = total
    return best, examined

tree_a = TreeNode(3,
                  TreeNode(2, None, TreeNode(3)),
                  TreeNode(3, None, TreeNode(1)))
tree_c = TreeNode(1, TreeNode(100), TreeNode(2, None, TreeNode(100)))
for name, tree in [("Tree A", tree_a), ("Tree C", tree_c)]:
    best, examined = brute(tree)
    print(name, "-> best haul:", best, "subsets examined:", examined)

Both answers match the cases, 7 and 200. Tree C has only four houses and already needs sixteen subsets. Twenty houses would need over a million.

The optimal solution: two numbers returned per node

The state at each node is a pair of numbers. The first is the best haul from that node's subtree when the node is robbed, and the second is the best haul when the node is skipped.

Robbing a node forces both children into their skipped numbers. Skipping a node lets each child choose its own better number. An absent child answers the pair (0, 0), since an empty subtree holds nothing either way.

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

def rob_pair(node):
    if node is None:
        return (0, 0)
    left = rob_pair(node.left)
    right = rob_pair(node.right)
    robbed = node.value + left[1] + right[1]
    skipped = max(left) + max(right)
    return (robbed, skipped)

tree_a = TreeNode(3,
                  TreeNode(2, None, TreeNode(3)),
                  TreeNode(3, None, TreeNode(1)))
tree_b = TreeNode(3,
                  TreeNode(4, TreeNode(1), TreeNode(3)),
                  TreeNode(5, None, TreeNode(1)))
tree_c = TreeNode(1, TreeNode(100), TreeNode(2, None, TreeNode(100)))
for name, tree, expected in [("Tree A", tree_a, 7), ("Tree B", tree_b, 9), ("Tree C", tree_c, 200), ("single house 5", TreeNode(5), 5)]:
    pair = rob_pair(tree)
    answer = max(pair)
    print(name, "root pair (robbed, skipped):", pair, "-> best haul:", answer, "match:", answer == expected)

All four cases match. Tree A does best by robbing its root, and Tree B does best by skipping its root. The pair keeps both options open until the root makes the final choice.

Walkthrough: Tree C from the leaves up

The walk computes Tree C bottom-up, naming each node. The leaves answer first. Their pairs flow into their parents, and the root's pair is read last.

class NamedNode:
    def __init__(self, name, value, left=None, right=None):
        self.name = name
        self.value = value
        self.left = left
        self.right = right

def rob_named(node):
    if node is None:
        return (0, 0)
    left = rob_named(node.left)
    right = rob_named(node.right)
    robbed = node.value + left[1] + right[1]
    skipped = max(left) + max(right)
    print("node", node.name, "- pair (robbed, skipped):", (robbed, skipped))
    return (robbed, skipped)

tree = NamedNode("root, value 1", 1,
                 NamedNode("left leaf", 100),
                 NamedNode("right child, value 2", 2, None, NamedNode("right leaf", 100)))
pair = rob_named(tree)
print("best haul:", max(pair))

The root's robbed number is 101, its own 1 plus the two skipped numbers below it. Its skipped number is 200, the two leaves taken through different branches. Skipping the root wins by a wide margin.

Complexity

The pair walk visits each node once and does fixed work per node, so its time is O(n). The recursion stack grows with the height of the tree, so the extra space is O(height).

The brute force examines 2 to the power n subsets. The counter below runs it on chains of houses, where the subset count is easy to read against the visit count of the pair walk.

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

def build_chain(values):
    root = None
    for value in reversed(values):
        root = TreeNode(value, None, root)
    return root

def brute_subsets(root):
    nodes = []

    def collect(node, parent):
        if node is None:
            return
        index = len(nodes)
        nodes.append((node.value, parent))
        collect(node.left, index)
        collect(node.right, index)

    collect(root, -1)
    examined = 0
    for mask in range(2 ** len(nodes)):
        examined += 1
    return examined

def pair_visits(node, counter):
    if node is None:
        return (0, 0)
    counter[0] += 1
    left = pair_visits(node.left, counter)
    right = pair_visits(node.right, counter)
    return (node.value + left[1] + right[1], max(left) + max(right))

for n in [5, 10, 15]:
    tree = build_chain([1] * n)
    counter = [0]
    pair_visits(tree, counter)
    print("houses", n, "- brute subsets:", brute_subsets(tree), " pair visits:", counter[0])

At 15 houses the subset count is 32,768 against 15 visits. Doubling backtracking is replaced by one bottom-up pass. Each node's pair summarises its whole subtree, so no choice is ever re-examined.

Edge cases and pitfalls

An empty tree holds nothing and answers 0. A single house answers its own value. A chain behaves like the street version of the problem, alternating rob and skip down its length.

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

def rob_pair(node):
    if node is None:
        return (0, 0)
    left = rob_pair(node.left)
    right = rob_pair(node.right)
    robbed = node.value + left[1] + right[1]
    skipped = max(left) + max(right)
    return (robbed, skipped)

print("empty tree pair:", rob_pair(None), "best:", max(rob_pair(None)))
print("single house:", rob_pair(TreeNode(5)))
chain = TreeNode(3, None, TreeNode(2, None, TreeNode(3)))
print("chain 3-2-3 root pair:", rob_pair(chain), "best:", max(rob_pair(chain)))
print("root skipped wins:", rob_pair(TreeNode(3, TreeNode(2), TreeNode(4))))
print("root robbed wins:", rob_pair(TreeNode(10, TreeNode(1), TreeNode(1))))

The chain row robs the root and the grandchild for 6. The pitfall in the pair formulas is swapping the two child numbers. Robbing takes the children's skipped numbers, never their better numbers, or neighbours end up robbed together.