Lab 08 — Tree DP (House Robber III)
Goal
Solve House Robber III (LC 337) with post-order DFS returning (rob, skip) per node. Internalize the post-order DP pattern where each node returns a tuple of “best with this node included” and “best with this node excluded”. After this lab you recognize tree DP within 60 seconds and can write any post-order tuple-DP from blank screen in <8 minutes.
A note on the four-stage progression: tree DP doesn’t have a clean tabulated form (there’s no natural row-major order for a tree), and “space optimization” is implicit (each post-order call returns a constant-size tuple, so the working memory is O(1) per node). Stages we can show: brute recursion, memoized recursion, post-order DFS with tuple return (the canonical form), and an iterative post-order using an explicit stack. The tuple-return version is what you write in interviews.
Background Concepts
Tree DP: state is per node; recurrence aggregates children’s states. The natural evaluation order is post-order — process all descendants before the node itself. Most tree DPs return a tuple (or struct) per node carrying the answers for “this node included” vs “this node excluded” (or whatever the binary split is). The parent combines children’s tuples in O(1) per child, giving O(N) total.
House Robber III is the canonical example. Each node v returns (rob_v, skip_v):
rob_v = val[v] + sum(skip_c for c in children(v))— robv, must skip every child.skip_v = sum(max(rob_c, skip_c) for c in children(v))— skipv, each child is independently best.
Final answer: max(rob_root, skip_root).
Interview Context
LC 337 is a top-30 Medium DP problem at Amazon and Microsoft. The post-order tuple pattern recurs in: LC 124 (Binary Tree Maximum Path Sum), LC 543 (Diameter of Binary Tree), LC 968 (Binary Tree Cameras), LC 1372 (Longest ZigZag Path). Mastering it here generalizes broadly. Senior interviewers specifically test whether you write the tuple-return version (clean, O(N)) versus the memoized-but-redundant version that recursively descends three times per call.
Problem Statement
A binary tree where each node holds an integer amount of money. The thief cannot rob two directly-linked houses (parent–child). Return the maximum amount the thief can rob without alerting the police.
Constraints
- 1 ≤ tree size ≤ 10^4
- 0 ≤
node.val≤ 10^4
Clarifying Questions
- Is the tree binary or general? (Binary, per LC 337.)
- Are values non-negative? (Yes — given.)
- What does the tree representation look like? (Standard
TreeNodewithleft,right.) - Can the tree be empty? (Yes — return 0.)
- Does “linked” mean parent–child only or also siblings? (Parent–child only.)
Examples
3
/ \
2 3
\ \
3 1 → 7 (rob 3 + 3 + 1)
3
/ \
4 5
/ \ \
1 3 1 → 9 (rob 4 + 5)
Initial Brute Force
For each subtree rooted at v: try rob-v (must skip children, recurse on grandchildren) or skip-v (recurse on children).
def rob_brute(root):
def f(v):
if v is None: return 0
# Skip v
skip = f(v.left) + f(v.right)
# Rob v: must skip both children
rob = v.val
if v.left: rob += f(v.left.left) + f(v.left.right)
if v.right: rob += f(v.right.left) + f(v.right.right)
return max(rob, skip)
return f(root)
Brute Force Complexity
Each call recurses on children and on grandchildren — the same node is visited multiple times via different paths. Worst case O(N · 2^depth). Memoization on the node identity collapses to O(N), but cleaner is to return the tuple in a single post-order traversal.
Optimization Path
The brute force descends three times per node (for the rob branch) and two for skip. With memoization on the node, every subtree is computed twice (once as a “child” call, once as a “grandchild” call). Use a dict keyed by id(node) or, much cleaner, return both values in one tuple per node — a single post-order pass with no memoization needed.
Final Expected Approach
Post-order DFS returning (rob, skip):
def f(v):
if v is None: return (0, 0)
lr, ls = f(v.left)
rr, rs = f(v.right)
rob_v = v.val + ls + rs
skip_v = max(lr, ls) + max(rr, rs)
return (rob_v, skip_v)
answer = max(f(root))
Time O(N) (each node visited once). Space O(H) for the call stack (H = tree height; O(N) worst case for skewed trees, O(log N) average).
Data Structures Used
- Binary tree (input).
- Recursion stack.
- For brute / memo: optional
dictkeyed by node identity.
Correctness Argument
By structural induction on the tree. Base: empty tree → (0, 0). Inductive: assume f(v.left) and f(v.right) correctly return (rob, skip) for those subtrees. Then:
rob_v= robvand skip both children. Since the children must be skipped (parent-child constraint), the contribution from each child subtree isskip_child. Plusv.val.skip_v= skipv, each child subtree independently maximized:max(rob_child, skip_child).
The max of the two is the overall optimum, but we return both (not the max) because the parent of v needs skip_v distinct from rob_v. The final answer at the root is max(rob_root, skip_root).
Complexity
| Stage | Time | Space |
|---|---|---|
| Brute force | O(N · 2^H) worst | O(H) |
| Memoized (node-keyed) | O(N) | O(N) |
| Tuple-return post-order | O(N) | O(H) |
| Iterative post-order | O(N) | O(N) (explicit stack) |
Implementation Requirements
class TreeNode:
def __init__(self, val=0, left=None, right=None):
self.val, self.left, self.right = val, left, right
# ---- Stage 1: Brute force (recompute via grandchildren) ----
def rob_brute(root):
def f(v):
if v is None: return 0
skip = f(v.left) + f(v.right)
rob = v.val
if v.left: rob += f(v.left.left) + f(v.left.right)
if v.right: rob += f(v.right.left) + f(v.right.right)
return max(rob, skip)
return f(root)
# ---- Stage 2: Memoized on node identity ----
def rob_memo(root):
memo = {}
def f(v):
if v is None: return 0
if id(v) in memo: return memo[id(v)]
skip = f(v.left) + f(v.right)
rob = v.val
if v.left: rob += f(v.left.left) + f(v.left.right)
if v.right: rob += f(v.right.left) + f(v.right.right)
memo[id(v)] = max(rob, skip)
return memo[id(v)]
return f(root)
# ---- Stage 3: Tuple-return post-order (canonical) ----
def rob(root):
def f(v):
if v is None: return (0, 0)
lr, ls = f(v.left)
rr, rs = f(v.right)
rob_v = v.val + ls + rs
skip_v = max(lr, ls) + max(rr, rs)
return (rob_v, skip_v)
return max(f(root))
# ---- Stage 4: Iterative post-order with explicit stack ----
def rob_iter(root):
if root is None: return 0
stack, order = [root], []
while stack:
v = stack.pop()
order.append(v)
if v.left: stack.append(v.left)
if v.right: stack.append(v.right)
# order is now reverse post-order; iterate reversed for true post-order
state = {} # id(v) -> (rob, skip)
for v in reversed(order):
lr, ls = state.get(id(v.left), (0, 0)) if v.left else (0, 0)
rr, rs = state.get(id(v.right), (0, 0)) if v.right else (0, 0)
state[id(v)] = (v.val + ls + rs, max(lr, ls) + max(rr, rs))
return max(state[id(root)])
Tests
Build trees from level-order input:
[3,2,3,null,3,null,1]→ 7.[3,4,5,1,3,null,1]→ 9.[1]→ 1.[]→ 0 (empty tree).- Linear left-skewed chain
1→2→3→4→5→ 9 (rob 1, 3, 5). - Random N=1000 tree — performance test.
- Cross-check all four implementations on random trees of size N≤8.
Follow-up Questions
- “Reconstruct which nodes were robbed.” → Augment the tuple to also return the set of robbed nodes (or backtrack from the root after the post-order pass).
- “K-ary tree.” → Same idea; sum over all children.
- “Constraint relaxes to: cannot rob two nodes within distance K.” → Augment state to track distance to last robbed node; state grows by factor K.
- “Negative values allowed.” → Same recurrence; max correctly handles.
- “Maximum path sum of an arbitrary path (LC 124).” → Different but related: post-order returns “best one-sided path from this node”. Combine left + right + node at the root for the best path through it.
Product Extension
Tree DP underlies code-review priority computation in commit-trees, expression-tree evaluation in compilers, hierarchical resource allocation (org charts, file systems), and game-tree value computation. The post-order-with-tuple pattern is a workhorse.
Language/Runtime Follow-ups
- Python: tuple unpacking is idiomatic. Default recursion limit (1000) overflows on deep skewed trees; bump with
sys.setrecursionlimit. - Java: define a small inner
int[]of size 2 or aPair<Integer, Integer>(or just useint[]). Recursion depth is O(H); JVM stack default 512KB, OK for N≤10000. - Go: return
(int, int)directly via multiple-return. - C++:
pair<int,int>returned by value. - JS/TS: return
[rob, skip]array.
Common Bugs
- Returning only
max(rob_v, skip_v)instead of the tuple — the parent then can’t distinguish the two cases and the recurrence is wrong. - Computing
rob_vasv.val + skip_v— wrong, becauseskip_valready includesmax(rob_child, skip_child)not justskip_child. - Forgetting
v is Nonebase case — null pointer /AttributeError. - Confusing which children’s value to sum:
rob_vsumsls + rs(skip both children);skip_vsumsmax(lr, ls) + max(rr, rs). - Iterative version: traversing in pre-order and reversing — works for binary trees because reverse pre-order with right-before-left equals post-order. Easy to flip and break.
- Using
@lru_cacheon the functionf(v)directly —TreeNodeis unhashable by default; useid(v)or define__hash__.
Debugging Strategy
For [3,2,3,null,3,null,1]:
- Leaves:
f(3-leaf-left) = (3, 0).f(1-leaf) = (1, 0).f(3-leaf-mid) = (3, 0). - Mid-left node (val 2, child 3):
f(2-mid) = (2 + 0, 3) = (2, 3). - Mid-right (val 3, right child 1):
f(3-right) = (3 + 0, 1) = (3, 1). - Root (val 3):
f(root) = (3 + 3 + 1, max(2,3) + max(3,1)) = (7, 6). Answer: 7. ✓
If your tuple values diverge, print (rob, skip) per node in post-order and locate the first inconsistency.
Mastery Criteria
- Recognized this as tree DP within 60 seconds.
-
Articulated the
(rob, skip)tuple invariant in <30 seconds. - Wrote the tuple-return post-order in <5 minutes from blank screen.
- Stated O(N) time and O(H) space.
- Articulated why the tuple is necessary (parent needs both rob_child and skip_child) in <30 seconds.
- Solved LC 337 unaided in <8 minutes.
- Solved LC 124 (Binary Tree Max Path Sum) in <12 minutes using the same post-order pattern.
- Solved LC 543 (Diameter) in <8 minutes.
- Identified the brute-force redundancy (descend twice via children + grandchildren) without prompting.