Tree DP and Rerooting: Computing an Answer at Every Node

A tree has no cycles, so the state graph of a tree DP is acyclic by construction and the only question is the traversal order. Root the tree arbitrarily, let each subtree be a subproblem, and the recursion runs from leaves upward: a node’s answer is a function of its children’s answers. That is the whole first pass, and it is O(n) for a state of constant size.

The difficulty begins with the second kind of question. Many problems ask for an answer at every node — the sum of distances from each node to all others, the eccentricity of each node, the number of nodes closer to v than to any other. Running the first pass once per node gives O(n²). Rerooting computes all of them in O(n) by running a second pass downward, and the technique requires the per-node combination to be invertible or to be handled with prefix and suffix aggregates.

Rooting and post-order

Pick any node as root. Orient every edge away from it. Each node v has a parent p(v) and a set of children c_1 … c_k. The subtree T_v is v together with the subtrees of its children, and the subproblems are exactly these subtrees — n of them, one per node, which is why tree DP is O(n) rather than O(n²) when the per-node state is small.

Post-order is the topological order: children before parents. It is the tree-shaped case of the ordering requirement discussed in recurrences as DAGs, and it is the only thing tree DP adds to ordinary tabulation. An explicit stack is safer than recursion, since a path graph of 10^6 nodes is a tree of depth 10^6 and will exhaust any default thread stack.

def subtree_sizes(adj, root):
    n = len(adj)
    parent = [-1] * n
    order = []
    stack = [root]
    parent[root] = root
    while stack:                       # build a rooted order, parents first
        v = stack.pop()
        order.append(v)
        for u in adj[v]:
            if parent[u] == -1:
                parent[u] = v
                stack.append(u)
    size = [1] * n
    for v in reversed(order):          # children before parents
        for u in adj[v]:
            if u != parent[v]:
                size[v] += size[u]
    return parent, size

The parent[u] == -1 test replaces a visited set: in a tree, the only already-visited neighbour of v is its parent, so marking parents is sufficient. In a general graph the same loop would revisit nodes and loop forever, which is the difference between tree DP and graph DP.

Merging children is a knapsack

When the state per node is not a scalar but a table over a resource — the number of selected nodes, the budget consumed, the height — the node combines its children’s tables by a convolution. Suppose v has children with subtree sizes s_1, …, s_k. Convolving the accumulated table, of width 1 + s_1 + … + s_{i-1}, with child i’s table of width s_i costs the product of the two widths. If the accumulated table is allocated at its final width n from the start, every merge costs O(n · s_i) and the node costs O(n · Σ s_i) = O(n²). Summed over the whole tree that is O(n³) in the worst case, which is a factor of n worse than necessary.

The standard improvement is to size the accumulated table at 1 + s_1 + … + s_i after each merge, so the convolution is over the actual merged width rather than n. Merging the first i children then costs (1 + s_1 + … + s_i) · s_{i+1}, and summing over i gives O((Σ s_i)²) for the subtree. Summed over the whole tree the total is O(n²), because every node appears in the subtree of each of its ancestors and the sum is bounded by counting each pair of nodes once at their lowest common ancestor. The structure is the interval-style merge enumeration of interval DP, applied over a node’s children instead of over a contiguous range.

The state itself has to be justified by the sufficiency rule in state sufficiency: a subtree is a complete summary of everything outside it only because the tree has no cycles, so the only route between the subtree and the rest of the graph is the single edge to the parent. On a general graph the same state is insufficient, and the repair is a boundary set — the treewidth formulation.

Rerooting: the answer for the tree rooted elsewhere

The first pass computes, for each node, the answer restricted to its subtree. The second pass computes the contribution of everything outside the subtree, and combines the two.

Define down[v] as the answer for T_v with v as root, computed bottom-up. Define up[v] as the answer for the whole tree with v as root, but computed from up[p(v)] and the siblings of v. Then the final answer at v combines down[v] with the contribution of up[v]’s side.

For the sum-of-distances problem the combination is simple because the operation is addition:

down[v] = Σ over children u of ( down[u] + size[u] )
up[v]   = (n - size[v]) + Σ over siblings u of ( down[u] + size[u] )

Read up[v] as “the sum of distances from v to every node outside T_v”. The first term accounts for the n − size[v] nodes outside T_v, each of which is one edge further from v than from p. The second term is the siblings’ contributions, which are exactly the parts of down[p] that do not come from v. The total at v is down[v] + up[v], and the answer at the root is just down[root] since up is empty there.

Subtracting down[v] + size[v] from down[p] recovers that sum in O(1), and the whole pass is O(n). This is the invertible case, and it is why the sum-of-distances rerooting is the standard worked example.

When the combination is not invertible

The subtraction fails for operations without inverses: maximum, minimum, greatest common divisor, bitwise OR, the second-largest value. For a max-combination, down[p] is the max over children, and “the max excluding v” cannot be recovered from the max including v.

The repair is prefix and suffix aggregates over the children’s contributions. For a node with children c_1 … c_k, compute pre[i] as the combination of c_1 … c_i and suf[i] as the combination of c_i … c_k; then the contribution excluding c_i is pre[i-1] ⊕ suf[i+1]. Building both arrays costs O(k) and answering all k exclusions costs O(k), so the pass stays linear in the number of edges. Storage is O(k) per node transiently, and the arrays can be reused per node if the traversal is iterative.

Two details matter in the implementation. The identity element for the combination must be defined — 0 for sum, −∞ for max, 1 for a product — and the prefix array must be initialized with it, or the first exclusion is wrong. And the operation must be one where “excluding one contributor” is well-defined: for a maximum attained by two children, excluding one leaves the maximum unchanged, which is the right answer for a maximum query but the wrong answer for a query that counts how many children attain it. Counts are not invertible in this sense, and prefix/suffix aggregation over the pairs (value, count) is the sound alternative.

For the maximum, a second scan over the same degree-sized list answers every exclusion, which is O(deg) per node and O(n) overall, and needs no prefix array at all. Keeping only the best value and the identity of its owner answers each exclusion in O(1), but only if ties are broken by a rule that survives them: the naive top-two scan must not attribute a value to the wrong neighbour when two neighbours tie. For a combination that is neither invertible nor summarizable by a small number of extremes, the prefix/suffix arrays are the general solution.

def farthest_down(adj, parent, order):
    down = [0] * len(adj)
    for v in reversed(order):
        best = 0
        for u in adj[v]:
            if u != parent[v]:
                best = max(best, down[u] + 1)
        down[v] = best
    return down

def farthest_all(adj, parent, order, down):
    n = len(adj)
    up = [0] * n
    for v in order:                       # parents before children
        # contribution of each neighbour: the distance it sits at, seen from v
        contrib = []
        for u in adj[v]:
            if u == parent[v]:
                contrib.append((up[v], u))     # the outside world, seen from v
            else:
                contrib.append((down[u] + 1, u))
        # two passes over a degree-sized list keep the worst case O(n)
        for u in adj[v]:
            if u != parent[v]:
                best = 0
                for val, w in contrib:
                    if w != u and val > best:
                        best = val
                up[u] = 1 + best
    return [max(down[v], up[v]) for v in range(n)]

The contrib list is rebuilt at every node and discarded after its children are assigned, so the transient storage is O(deg(v)) and the total work is proportional to the number of edges.

What rerooting actually assumes

The technique rests on a decomposability property: the answer at v is a combination of one contribution per incident edge, and each edge’s contribution depends only on the component on the far side. This holds for path-based quantities — distances, eccentricities, subtree sizes, counts of nodes within a radius — because deleting v splits the tree into exactly deg(v) components, and the route from v to any other node leaves through exactly one of them.

It fails when a quantity couples two components at once. The number of pairs of nodes at distance k, for a fixed k, decomposes into pairs inside a single component plus pairs split across two components whose distances sum to k - 2. The crossing term is a convolution over the components, and the per-node work becomes O(deg · k) rather than O(deg) — the same shape as the child-merging convolution above, with the distance bound playing the role of the table width. Prefix and suffix aggregation still applies, but over tables rather than scalars.

A third category is quantities that do not depend on the root at all. The tree’s diameter is the same number whichever node is designated root, so rerooting it computes one value many times. The root- dependent quantity worth computing instead is the eccentricity: the maximum distance from v to any other node, which does decompose per incident edge because the route to a node passes through exactly one of them. Distinguishing root-independent quantities from root-dependent ones is the first step in deciding whether a second pass is needed at all.

Common failure modes

Depth. Recursion on a path graph is O(n) deep. An iterative post-order with an explicit stack, as in the first code block, removes the limit at the cost of a heap allocation for the order array.

Parent handling in the second pass. The up array must be filled in pre-order — parents before children — and the value up[v] used for the outside contribution must be the value for v, not the value for p(v). Mixing the two is the most common off-by-one-level error, and it produces answers that are correct for the root and wrong elsewhere.

Combining up and down inconsistently. The final answer at v combines down[v], which covers T_v, with up[v], which covers the complement. If either includes the node v itself twice, or neither does, the result is off by a constant per node — a bug that passes tests on a two-node tree and fails on larger ones.

Order of the two passes. The invariant is that when up[u] is assigned for a child u, three things are already final: up[v] for its parent, and the down values of every child of v, including the siblings of u. The two-loop form above satisfies this by construction. A single recursive traversal also satisfies it, provided up[u] is assigned after the whole node’s post-order work is complete and before descending into u’s own second-pass work. Assigning up[u] before the siblings’ down values are known breaks the invariant, and the resulting answers are correct for the root and its immediate children and wrong below that.

The technique generalizes beyond trees to any graph of bounded treewidth, where the DP state is a boundary set rather than a single node and the table is exponential in the width. At width one — a tree — the state is constant-sized and the whole computation is linear; each unit of additional width multiplies the table by the size of the state alphabet.