Tree DP, in depth: designing subtree states, iterative evaluation, O(nK) knapsack merges, rerooting and brute-force testing

The recipe: state, combine, order

Root the tree anywhere. For each vertex v define a small vector dp[v] that summarises everything the rest of the tree needs to know about v's subtree. The key design question is: what does my parent need to know about me to make its own decision? If the parent only needs the best value, one number is enough. If the parent's choice depends on whether v is selected, covered, coloured red, or holds j chosen vertices, then dp[v] needs one entry per such case.

Three things then define the algorithm:

1. The state. One entry per situation at v that the parent can tell apart. Too few entries and the recurrence is wrong; too many and it is slow.
2. The combine. How dp[v] is built from dp[c] for each child c. It must look only at children, never at grandchildren or siblings' internals.
3. The order. Children before parents. Any post-order works; the answer is read at the root.

The classic first example, the maximum-weight independent set with two states (v taken, v not taken), is covered with reconstruction in the house robber article, and the longest-path version with its top-two trick in tree diameter in depth. Here we go one level harder.

Evaluating without recursion

The textbook version is a recursive DFS. Python's default recursion limit is 1000, so a path-shaped tree of 100,000 vertices crashes it, and raising the limit moves the crash into the C stack. Instead compute a BFS order once: every parent appears before its children, so iterating it reversed visits children before parents, which is all a bottom-up DP needs, and iterating it forwards gives the top-down order rerooting needs.

from collections import deque

def root_tree(n, edges, root=0):
    """Return parent, BFS order and children lists for a tree given as an edge list."""
    adj = [[] for _ in range(n)]
    for a, b in edges:
        adj[a].append(b)
        adj[b].append(a)
    parent = [-1] * n
    seen = [False] * n
    seen[root] = True
    order, q = [], deque([root])
    while q:
        v = q.popleft()
        order.append(v)
        for u in adj[v]:
            if not seen[u]:
                seen[u] = True
                parent[u] = v
                q.append(u)
    if len(order) != n:
        raise ValueError("graph is not connected: this is a forest, root each component")
    children = [[] for _ in range(n)]
    for v in order[1:]:
        children[parent[v]].append(v)
    return parent, order, children

The connectivity check matters: a forest silently produces dp values only for the root's component, and the missing vertices look like zeros. For the explicit-stack alternative and its pitfalls, see iterative DFS in depth.

Worked example: minimum dominating set with three states

A dominating set is a set S of vertices such that every vertex is in S or adjacent to a vertex in S: think of placing the fewest monitoring agents so every server is watched by itself or a neighbour. On general graphs the problem is NP-hard. On trees it is linear, and it is the right example for state design because two states are not enough.

What does a parent need to know about child c? Whether c is in S (then c also covers the parent), whether c is already covered by one of its own children, or whether c is still uncovered and is relying on the parent to join S. That gives three states:

• IN: v is in S; every vertex in v's subtree is covered.
• COVERED: v is not in S but at least one child is; every vertex in the subtree is covered.
• NEED: v is not in S, no child is in S, so v is uncovered; everything below v is covered. Only legal if the parent joins S.

The transitions follow directly. If v is IN, each child may be in any state, because v covers a NEED child. If v is NEED, every child must be COVERED: a child in S would cover v, contradicting the state, and a NEED child would stay uncovered. COVERED is the subtle one: every child is IN or COVERED, and at least one must be IN. Take the cheaper of the two for each child, then pay the smallest extra cost to force one child to IN.

INF = float("inf")

def min_dominating_set(n, edges):
    parent, order, children = root_tree(n, edges)
    dp = [None] * n                      # (IN, COVERED, NEED)
    for v in reversed(order):
        take, free, bump, need = 1, 0, INF, 0
        for c in children[v]:
            i, cov, nd = dp[c]
            take += min(i, cov, nd)
            best = min(i, cov)
            free += best
            bump = min(bump, i - best)   # extra cost to force this child into S
            need += cov
        dp[v] = (take, free + bump, need)
    i, cov, _ = dp[order[0]]
    return min(i, cov)                   # the root has no parent to rescue NEED

assert min_dominating_set(4, [(0, 1), (1, 2), (2, 3)]) == 2
assert min_dominating_set(5, [(0, 1), (0, 2), (0, 3), (0, 4)]) == 1

Trace the path 0-1-2-3 rooted at 0. Leaf 3 is (1, INF, 0): it can join S, it cannot be covered by children it does not have, or it can wait for its parent. Vertex 2 becomes (1, 1, INF), vertex 1 becomes (2, 1, 1) and the root becomes (2, 2, 1). The answer is min(2, 2) = 2, for example S = {1, 2}. Notice that NEED at the root is excluded, the most common bug in this problem.

Knapsack on a tree and the size-clamp bound

Many tree problems carry a budget: choose at most K vertices, where a vertex can only be chosen if its parent is, and maximise total value, as when picking modules of a dependency tree. Now dp[v][j] is the best value using exactly j chosen vertices from v's subtree, with v itself chosen. Combining a child is a knapsack merge: new[a + b] = max(dp[v][a] + dp[c][b]).

A naive merge loops a and b each over 0..K for every child, which is O(n·K²). The fix is one line of discipline: loop each side only up to the size of the subtree merged so far, clamped to K. Without the clamp, the work for merging sizes s₁ and s₂ is s₁·s₂, which you can charge to pairs of vertices that meet for the first time at their lowest common ancestor; each pair meets once, so the total is O(n²). With the clamp the same charging argument gives O(n·K). For n = 10⁵ and K = 100 that is the difference between 10⁷ and 10⁹ inner-loop steps.

NEG = float("-inf")

def tree_knapsack(n, edges, value, K, root=0):
    """Best total value of a parent-closed set of at most K vertices containing the root."""
    parent, order, children = root_tree(n, edges, root)
    dp, size = [None] * n, [0] * n
    for v in reversed(order):
        cur, sz = [NEG, value[v]], 1          # cur[j]: j chosen in merged part, v chosen
        for ch in children[v]:
            child, csz = dp[ch], size[ch]
            new = [NEG] * (min(sz + csz, K) + 1)
            for a in range(1, min(sz, K) + 1):
                if cur[a] == NEG:
                    continue
                if cur[a] > new[a]:
                    new[a] = cur[a]           # take nothing from this child
                for b in range(1, min(csz, K - a) + 1):
                    s = cur[a] + child[b]
                    if s > new[a + b]:
                        new[a + b] = s
            cur, sz = new, sz + csz
            dp[ch] = None                     # free the child's table
        dp[v], size[v] = cur[:K + 1], sz
    return max(dp[root][1:])

# root 0 (value 1) with children 1 (5) and 2 (2); vertex 3 (10) hangs under 2
assert tree_knapsack(4, [(0, 1), (0, 2), (2, 3)], [1, 5, 2, 10], K=3) == 13

With K = 3 the best set is {0, 2, 3} worth 13, beating {0, 1, 2} worth 8: the high-value leaf is only reachable by paying for its parent, which is precisely the dependency the DP encodes. Releasing child tables after the merge keeps peak memory at O(n) rather than O(n·K): live tables belong to disjoint subtrees and each is at most min(size, K) + 1 long.

Rerooting: every vertex's answer in two passes

Some questions ask for an answer at every vertex: the sum of distances from each node to all others, or the farthest node from each server. Running the DP from every root costs O(n²). Rerooting gets all n answers in O(n) with two passes. The first pass computes subtree values for one root. The second walks top-down and asks: if the root moved from parent p to child v, how does the answer change?

For the sum of distances the change is clean. Moving the root across edge (p, v) brings the size[v] vertices of v's subtree one step closer and pushes the other n - size[v] vertices one step further away, so ans[v] = ans[p] - size[v] + (n - size[v]).

def sum_of_distances(n, edges):
    parent, order, _ = root_tree(n, edges)
    size, down = [1] * n, [0] * n
    for v in reversed(order):                 # pass 1: children before parents
        p = parent[v]
        if p >= 0:
            size[p] += size[v]
            down[p] += down[v] + size[v]
    ans = [0] * n
    ans[order[0]] = down[order[0]]
    for v in order[1:]:                       # pass 2: parents before children
        ans[v] = ans[parent[v]] - size[v] + (n - size[v])
    return ans

assert sum_of_distances(6, [(0, 1), (0, 2), (2, 3), (2, 4), (2, 5)]) == [8, 12, 6, 10, 10, 10]

Check vertex 2 by hand: it is 1 step from 0, 3, 4 and 5 and 2 steps from 1, total 6. When the combine is not invertible, as with max, you cannot subtract a child's contribution back out. Instead, for each vertex keep prefix and suffix combinations over its children, so the value "everything except child i" costs O(1). The LCA article applies the same move to path queries.

Failure modes

These are the bugs that show up in reviews and wrong-answer verdicts, roughly in order of frequency.

• Missing state. A two-state dominating-set DP passes small tests and fails on paths, because it cannot express "uncovered, waiting for my parent". Write down what the parent must distinguish before writing code.
• Reading an illegal root state. NEED at the root must be excluded at the end.
• Recursion depth. Recursive DFS dies on deep trees in Python and can overflow native stacks in C++. Use the BFS order.
• Unclamped merges. Looping to K instead of min(size, K) silently turns O(n·K) into O(n·K²); tests on small trees never notice.
• Using a stale table. Merging into the same array you read from double-counts a child. Write into a fresh array, or iterate a downwards as in 0/1 knapsack.
• Overflow and sentinels. INF arithmetic is safe in Python but wraps in fixed-width integers; use a sentinel like 10¹⁸ with explicit checks.
• Forests and bad input. Verify n - 1 edges and full connectivity once at the start.

Testing against brute force

Tree DP bugs have few visible symptoms, so test against brute force: random trees of 1 to 10 vertices built with parent[i] = a random vertex below i, which yields paths, stars and everything between, checked by enumerating subsets. A thousand trees take a second and catch every bug above except overflow.

import itertools, random

def brute_dominating(n, edges):
    nbr = [{v} for v in range(n)]
    for a, b in edges:
        nbr[a].add(b); nbr[b].add(a)
    for k in range(1, n + 1):
        for S in itertools.combinations(range(n), k):
            if all(nbr[v] & set(S) for v in range(n)):
                return k

for _ in range(1000):
    n = random.randint(1, 10)
    edges = [(i, random.randrange(i)) for i in range(1, n)]
    assert min_dominating_set(n, edges) == brute_dominating(n, edges)

Trade-offs

What to do next

1. Type in root_tree and min_dominating_set, run the asserts and the brute-force harness, then delete the NEED exclusion at the root and watch the harness catch it.
2. Remove the clamp in tree_knapsack, time both versions on a random 20,000-vertex tree with K = 200, and confirm the gap matches the O(n·K) argument.
3. Extend sum_of_distances to weighted edges: the shift becomes w(p, v) times (n - 2·size[v]).
4. Solve the farthest-node-from-every-vertex problem with prefix and suffix maxima, and check it against n BFS runs on random trees.
5. Write a one-line comment above every DP stating what each state means: it is the first thing a reviewer needs.