Tree DP: Subtree States, Path Sums, Tree Knapsack and Rerooting

Tree DP computes an answer for every subtree from the answers of its children. A tree has no cycles, so once it is rooted, the subtrees form a natural DAG of subproblems and a post-order traversal visits them in a valid order.

Intuition

Root the tree anywhere. For a node u, "the best answer inside u's subtree" usually depends only on the best answers inside each child's subtree, plus a small amount of information about how u itself is used. That extra information becomes part of the state:

• take[u] / skip[u]: is u chosen or not (independent sets, robbers, vertex covers).
• "best path going down from u" versus "best path bending at u" (path sums, diameters).
• f[u][j]: best result using j nodes of the subtree (tree knapsack).

Compute children first, then combine. When the answer is needed for every possible root, a second top-down pass ("rerooting") reuses the first pass instead of repeating it n times.

When to use it in interviews / OAs

• The input is a tree (or a hierarchy: managers and employees, prerequisites, file systems) and you optimize or count something over it.
• Adjacent nodes constrain each other ("cannot pick a node and its parent", "every edge must be covered").
• A path problem on a tree where the best path may bend at any node.
• "For every node, compute ... as if it were the root" means rerooting.

Template 1: take or skip a node

Pick nodes to maximize total weight so that no picked node is adjacent to another picked node (an employee does not come if their direct manager comes). If u is taken, its children must be skipped; if u is skipped, each child does whatever is better.

def max_independent_set(n: int, edges: list[tuple[int, int]], weight: list[int], root: int = 0) -> int:
    adj = [[] for _ in range(n)]
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)
    parent, order = [-1] * n, [root]
    for u in order:  # BFS order: every parent appears before its children
        for v in adj[u]:
            if v != parent[u]:
                parent[v] = u
                order.append(v)
    skip = [0] * n  # best total in u's subtree when u is not picked
    take = [0] * n  # best total in u's subtree when u is picked
    for u in reversed(order):  # children are finished before their parent
        take[u] += weight[u]
        if u != root:
            p = parent[u]
            skip[p] += max(skip[u], take[u])
            take[p] += skip[u]
    return max(skip[root], take[root])

Processing nodes in reverse BFS order is an iterative post-order, so a path-shaped tree with 10^5 nodes does not hit Python's recursion limit. On a binary tree given as TreeNode objects, the same logic is a short recursive function returning the pair (take, skip).

Template 2: return one thing, record another

The maximum path sum in a binary tree: a path may bend at one node and use both of its children, but what a node can pass up to its parent is only a straight downward path. So the DFS returns the best downward path and records the best bent path seen anywhere.

class TreeNode:
    def __init__(self, val: int = 0, left: "TreeNode | None" = None, right: "TreeNode | None" = None):
        self.val, self.left, self.right = val, left, right


def max_path_sum(root: TreeNode) -> int:
    best = float("-inf")

    def down(node: TreeNode | None) -> int:
        # best sum of a path that starts at node and only goes downward
        nonlocal best
        if node is None:
            return 0
        left = max(down(node.left), 0)  # a negative branch is better left out
        right = max(down(node.right), 0)
        best = max(best, node.val + left + right)  # the path that bends at node
        return node.val + max(left, right)

    down(root)
    return best

The tree diameter, the longest path with distinct adjacent labels, and many "longest/heaviest path in a tree" questions use the same two-value pattern.

Template 3: tree knapsack

Courses have at most one prerequisite, which makes them a forest. You must take exactly m courses, a course only together with its prerequisite, and you want the maximum total score. Hang every root on a virtual course with score 0, and let f[j] for node u be the best score taking j courses from u's subtree, where j >= 1 means u itself is taken. Children are merged into the parent like a group knapsack.

def best_with_prerequisites(prereq: list[int], score: list[int], m: int) -> int:
    # prereq[i] = the course required before course i, or -1; assumes 1 <= m <= len(prereq)
    n = len(prereq)
    root = n  # virtual course with score 0 above every course that has no prerequisite
    children = [[] for _ in range(n + 1)]
    for i, p in enumerate(prereq):
        children[root if p == -1 else p].append(i)
    val = score + [0]
    NEG = float("-inf")

    def dfs(u: int) -> list[float]:
        # f[j] = best score taking j courses in u's subtree; j >= 1 means u is taken
        f = [0, val[u]]
        for v in children[u]:
            g = dfs(v)
            size = min(len(f) + len(g) - 1, m + 2)  # never need more than m + 1 nodes
            merged = f + [NEG] * (size - len(f))  # taking nothing from v
            for a in range(1, len(f)):
                for b in range(1, len(g)):
                    if a + b < size:
                        merged[a + b] = max(merged[a + b], f[a] + g[b])
            f = merged
        return f

    return int(dfs(root)[m + 1])  # the virtual root is one of the m + 1 taken nodes

Capping every array at the subtree size (and at m + 1) is what makes the total work O(n * m) instead of O(n * m^2).

Template 4: rerooting

Sum of distances from every node to all others. First pass, rooted at 0: subtree sizes, and the answer for node 0 (each non-root node u contributes size[u], the number of nodes below the edge to its parent). Second pass: moving the root from p to its child u makes size[u] nodes one step closer and n - size[u] nodes one step farther.

def sum_of_distances(n: int, edges: list[tuple[int, int]]) -> list[int]:
    adj = [[] for _ in range(n)]
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)
    parent, order = [-1] * n, [0]
    for u in order:
        for v in adj[u]:
            if v != parent[u]:
                parent[v] = u
                order.append(v)
    size, ans = [1] * n, [0] * n
    for u in reversed(order[1:]):  # pass 1: bottom-up sizes and the answer for root 0
        size[parent[u]] += size[u]
        ans[0] += size[u]
    for u in order[1:]:  # pass 2: top-down, move the root from parent[u] to u
        ans[u] = ans[parent[u]] + n - 2 * size[u]
    return ans

Complexity

Pattern	Time	Space
Take/skip, path sums	O(n)	O(n)
Tree knapsack with size caps	O(n * m)	O(n * m) in the worst case
Rerooting (two passes)	O(n)	O(n)

Common pitfalls

• Recursion depth: a path of 10^5 nodes overflows Python's default stack. Use the BFS-order trick above, or raise the recursion limit for small inputs.
• Undirected edges: always skip the parent when iterating neighbors, or you will walk back up.
• Returning the wrong quantity: the value passed up must be extendable by the parent (a downward path), while the answer may be something that cannot be extended (a bent path).
• Negative values: decide whether a branch can be dropped (max(child, 0)) or must be included.
• Tree knapsack loops: iterating up to m for every child instead of up to the current sizes costs an extra factor of m.
• Rerooting formula: derive it from what changes when the root moves across one edge; test it on a 3-node path.

Interview variations

• House robber on a tree, maximum independent set, employee party: take/skip states.
• Binary tree cameras or minimum vertex cover: three states per node (has camera, covered without camera, not covered).
• Maximum path sum, diameter, longest path with different adjacent characters: return a downward value, record a bent value.
• Distribute coins / balance a tree: each edge carries abs(excess of the subtree) moves.
• Sum of distances in a tree, minimum height trees: rerooting.
• Choose k nodes forming a connected group with a fixed root, courses with prerequisites: tree knapsack.