bigonotes
DSA: Patterns & PythonTrees

Tree DP with States

Return a small tuple of states from each subtree, combine them in postorder, and read the answer off the root.

6 min readLast updated

Tree DP is the bottom-up pattern with more than one answer per node. When a parent's choice depends on a choice made at its child (taken or not, camera or not), one number is not enough. Each subtree returns a small tuple, one value per state, and the parent combines its children's tuples in postorder. Every node is visited once, so it runs in O(n)O(n) with no memo table.

The idea#

Write down the states a node can be in, then write how each state is computed from the children's states. That is the whole recipe:

  1. States: what does the parent need to know about this subtree? For example, "best total if this node is taken" and "best total if it is skipped".
  2. Transitions: each state as a formula over the children's states. A rule such as "parent and child can't both be taken" decides which child states each formula may use.
  3. Base case: an empty subtree returns neutral values, such as (0, 0).
  4. Answer: combine the root's states, for example with max.

You may first think of memoizing on (node, flag), where the flag records the parent's choice (for example, "the parent was taken"). That version reaches the same answer and is a good stepping stone. Returning the tuple is cleaner: the postorder walk is the bottom-up DP, so there is no cache to manage and nothing is computed twice.

How to spot it#

The parent's options depend on what was decided at the child.

  • "No two adjacent nodes" (parent and child can't both be chosen).
  • Covering, monitoring or guarding nodes with the fewest picks.
  • Paths that can go up one side and down the other with a condition on each step (increasing on one side, decreasing on the other).
  • Moving units between nodes, where each edge carries the excess of a subtree.

Template#

Here is that recipe as code: one tuple of states per call, combined in postorder. Nodes follow LeetCode's TreeNode(val, left, right).

solve(node):                       # returns one value per state
    if node is None: return the neutral tuple, e.g. (0, 0)
    left = solve(node.left)
    right = solve(node.right)
    state_a = formula over node, left and right
    state_b = formula over node, left and right
    return (state_a, state_b)
answer = combine the states of solve(root)

Greedy with states: what do I owe, what can I offer#

The template keeps every state and lets the root pick. Some tree problems are greedy instead: they decide from the leaves up, with a small set of states. A parent's choice depends entirely on what its children report, so process children first. Each node sends its parent one message: what it still needs, and what it can give. Two details matter. The order in which you check the children's messages can change the result. The root has no parent to ask, so handle its message after the walk.

Both shapes push answers only from the leaves up. When every node needs the answer as if it were the root, there is a further step.

Worth exploring: rerooting DP (Sum of Distances in Tree, LC 834). A postorder pass computes each subtree's answer from the leaves up, then a preorder pass pushes answers from the root down, so every node gets the answer as if it were the root.

Pitfalls#

These mistakes lose information or double the work.

  • Taking the max too early. If a call returns only the best of its states, the parent can't tell which choice the child made. Return all states and take the max at the root.
  • A redundant outer max. In the memoized (node, flag) version, the "free to choose" state already takes the max. So taking the max of both root calls repeats work.
  • Check order in multi-state greedy. When two children send different messages, decide which message wins before you code.
  • Forgetting the root. A state that asks the parent for something needs special handling at the root.

Problems in this pattern#

These problems use this pattern. Solve them before you open the hints below. Start with the first one.

Hints#

Open a hint only after you have tried the problem for a while.

Hint: House Robber III

Return (rob_this, skip_this). Robbing forbids robbing the children: node.val + left.skip + right.skip. Skipping leaves each child free: max(left) + max(right). The answer is max of the root's pair.

Hint: Binary Tree Cameras

Use three states: 0 "cover me" (not covered), 1 "I'm fine, but I can't help you" (covered, no camera), 2 "I've got you" (has a camera). An empty child reports 1, so a leaf asks its parent for cover. Putting the camera on the leaf's parent is the greedy step: it also covers its own parent and its other child.

Hint: Binary Tree Cameras: check order

If one child needs cover and the other has a camera, the node must still buy a camera. So check "a child needs cover" before "a child has a camera". If the root ends up asking for cover, add one camera.

Hint: Binary Tree Longest Consecutive Sequence II

Return (inc, dec): the longest run ending at this node going up by 1, and going down by 1. Extend a child's run only when its value differs from the node's by exactly 1. The best path through the node is inc + dec - 1, because the node counts in both runs.

Hint: Binary Tree Maximum Path Sum

This is the one-state version: return the best downward gain and record the best bent path in an outer variable. See Bottom-Up Return Values.

Hint: Distribute Coins in Binary Tree

Return the excess of each subtree: its coins minus its nodes, which can be negative. Every unit of excess must cross the edge to the parent. So add abs(left) + abs(right) to the move count at each node.

Hint: Maximum Product of Splitted Binary Tree

Get the total sum in a first pass. In a second postorder pass, each subtree sum s is one possible cut, worth s * (total - s). Track the maximum and apply the modulo only at the end.

Key takeaways#

  • Tree DP is postorder with a tuple of states per node.
  • Name the states, write each one from the children's states, set the empty-tree base case, read the answer at the root.
  • Return all states; take the max only at the root.
  • Some tree DPs are greedy with states, like cameras: "what do I owe, what can I offer".
  • One pass is O(n)O(n) with no memo table.
Revision notes · 2 min

Core idea#

Tree DP is the bottom-up pattern with more than one answer per node. Each subtree returns a small tuple, one value per state. The parent combines the tuples in postorder in O(n)O(n), with no memo table.

Key points#

  • Recipe: name the states, write each from the children's states, set the empty base case, read the answer at the root.
  • House Robber III: rob = node.val + left.skip + right.skip. Skip = max(left) + max(right).
  • Return both states. Take the max only at the root.
  • A memoized (node, must_skip) version works, but the tuple has no cache to manage.
  • Increasing and decreasing runs: return (inc, dec). The path through the node is inc + dec - 1.
  • Spot it: no two adjacent nodes, fewest picks to cover nodes, units moving along edges.

Cameras: three states#

State Message to the parent
0 NEEDS_COVER "Cover me."
1 COVERED "I am fine, but I cannot help you."
2 HAS_CAMERA "I have got you."
  • An empty child reports COVERED, so a leaf asks its parent for cover.
  • A camera on the leaf's parent is the greedy step. It also covers its own parent and other child.

Template#

best(node):                          # returns (rob_this, skip_this)
    if node is None: return (0, 0)
    l_rob, l_skip = best(node.left)
    r_rob, r_skip = best(node.right)
    rob_this = node.val + l_skip + r_skip      # children skipped
    skip_this = max(l_rob, l_skip) + max(r_rob, r_skip)
    return (rob_this, skip_this)
answer = max of best(root)

Pitfalls#

  • Returning only max(rob, skip). The parent cannot tell if the child was robbed.
  • Wrong check order in cameras. Handle "a child needs cover" before "a child has a camera".
  • Forgetting the root. If the root needs cover, add one camera.

Say it in an interview#

  • "Each node returns the best total with and without it, and the parent combines them."
  • "It is one postorder pass, so O(n) with no memo."

Found a mistake? Email ask@bigonotes.com