DP on Trees: From Max Independent Set to Rerooting

Trees and dynamic programming go together like recursion and base cases. Every subtree is an independent subproblem, and the answer at a parent depends only on the answers at its children. That structure makes trees one of the cleanest settings for DP — no cycles, no overlapping paths to worry about, and the post-order traversal hands you the answers bottom-up.

This post covers the core tree DP patterns you will encounter in interviews and competitive programming, building from simple leaf-to-root problems up to the powerful rerooting technique.

---

1. The Core Idea — Post-Order DP

In tree DP, you compute the answer for each node after computing it for all its children. This is exactly post-order traversal. For each node u, you store one or more DP values that summarise the optimal answer for the subtree rooted at u.

def tree_dp(node, parent, adj):
    """Generic tree DP skeleton using adjacency list."""
    # dp[node] will hold the answer for the subtree rooted at node
    dp[node] = base_value

    for child in adj[node]:
        if child == parent:
            continue
        tree_dp(child, node, adj)
        # Combine child's answer into node's answer
        dp[node] = combine(dp[node], dp[child])

The key insight: because a tree has no cycles, visiting children before the parent guarantees that every subproblem is solved exactly once. The total work is O(n) for n nodes.

---

2. Max Independent Set on a Tree

Problem: given a tree where each node has a weight, pick a subset of nodes with maximum total weight such that no two picked nodes are adjacent.

This is the tree version of the classic “House Robber” problem. For each node, we track two states:

• dp[u][0] — the best answer if we do not include node u
• dp[u][1] — the best answer if we do include node u

Recurrence

• If we include u, we cannot include any child: dp[u][1] = weight[u] + sum(dp[c][0] for c in children)
• If we exclude u, each child can be included or not: dp[u][0] = sum(max(dp[c][0], dp[c][1]) for c in children)

Python Implementation

import sys
from collections import defaultdict

sys.setrecursionlimit(200_000)


def max_independent_set(n, edges, weights):
    """
    Find max-weight independent set on a tree.

    Args:
        n: number of nodes (0-indexed)
        edges: list of (u, v) pairs
        weights: list of weights for each node

    Returns:
        Maximum total weight of an independent set.
    """
    adj = defaultdict(list)
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)

    # dp[node] = (exclude_value, include_value)
    dp = [(0, 0)] * n

    def dfs(node, parent):
        include = weights[node]
        exclude = 0

        for child in adj[node]:
            if child == parent:
                continue
            dfs(child, node)
            child_exc, child_inc = dp[child]

            # If we include node, children must be excluded
            include += child_exc
            # If we exclude node, children can go either way
            exclude += max(child_exc, child_inc)

        dp[node] = (exclude, include)

    dfs(0, -1)
    return max(dp[0])


# Example
n = 7
edges = [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6)]
weights = [3, 4, 5, 3, 2, 1, 2]

print(max_independent_set(n, edges, weights))
# Output: 14 (nodes 1, 2, and leaves 3+2+1+2... pick 4+5+3+2 = 14)

Time: O(n) — each node visited once. Space: O(n) — DP array plus recursion stack.

---

3. House Robber III — Binary Tree Version

LeetCode 337 is the binary-tree specialisation. The tree is given as a TreeNode with .left and .right.

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


def rob(root: TreeNode) -> int:
    """
    Return the maximum amount you can rob without robbing
    two directly connected houses.
    """

    def dfs(node):
        if not node:
            return (0, 0)  # (exclude, include)

        left_exc, left_inc = dfs(node.left)
        right_exc, right_inc = dfs(node.right)

        # Include this node: children must be excluded
        include = node.val + left_exc + right_exc

        # Exclude this node: children can go either way
        exclude = max(left_exc, left_inc) + max(right_exc, right_inc)

        return (exclude, include)

    exc, inc = dfs(root)
    return max(exc, inc)


# Build example tree:
#       3
#      / \
#     4   5
#    / \   \
#   1   3   1
root = TreeNode(3)
root.left = TreeNode(4, TreeNode(1), TreeNode(3))
root.right = TreeNode(5, None, TreeNode(1))

print(rob(root))  # Output: 9 (rob 4 + 5 = 9? No — rob 4+5+1? Actually 1+3+5=9 or 4+1+5=... let's check: 4+5+1=10? no, 4 and 5 are children of 3 so they CAN be picked together since they're not adjacent to each other. 4+5+1=10? Actually 4+5 are siblings not adjacent. Pick 4+5+1(right leaf)=10.)

The structure is identical to the general tree version. The binary-tree form just makes the child iteration simpler.

---

4. Tree Diameter via DP

The diameter of a tree is the longest path between any two nodes (measured in edges or weights). A common approach is two BFS passes, but DP gives an elegant single-pass solution.

Idea: for each node u, compute the longest downward path from u into its subtree. The diameter through u is the sum of the two longest downward paths from u through different children.

def tree_diameter(n, edges):
    """
    Compute the diameter of an unweighted tree.

    Returns:
        The number of edges on the longest path.
    """
    adj = defaultdict(list)
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)

    diameter = 0

    def dfs(node, parent):
        nonlocal diameter
        # Longest and second-longest downward paths
        max1 = max2 = 0

        for child in adj[node]:
            if child == parent:
                continue
            child_depth = dfs(child, node) + 1

            if child_depth >= max1:
                max2 = max1
                max1 = child_depth
            elif child_depth > max2:
                max2 = child_depth

        # Diameter through this node
        diameter = max(diameter, max1 + max2)

        return max1  # longest downward path from this node

    dfs(0, -1)
    return diameter


# Example: path graph 0-1-2-3-4
edges = [(0, 1), (1, 2), (2, 3), (3, 4)]
print(tree_diameter(5, edges))  # Output: 4

Why it works: at every node we check whether the longest path in the whole tree passes through it. By tracking the two deepest subtrees, we cover all candidates in a single O(n) pass.

---

5. Weighted Tree Diameter

For weighted edges, the same pattern applies — just add edge weights instead of counting edges.

def weighted_tree_diameter(n, weighted_edges):
    """
    Compute diameter of a weighted tree.

    Args:
        weighted_edges: list of (u, v, w) tuples

    Returns:
        Maximum path weight between any two nodes.
    """
    adj = defaultdict(list)
    for u, v, w in weighted_edges:
        adj[u].append((v, w))
        adj[v].append((u, w))

    diameter = 0

    def dfs(node, parent):
        nonlocal diameter
        max1 = max2 = 0

        for child, weight in adj[node]:
            if child == parent:
                continue
            child_dist = dfs(child, node) + weight

            if child_dist >= max1:
                max2 = max1
                max1 = child_dist
            elif child_dist > max2:
                max2 = child_dist

        diameter = max(diameter, max1 + max2)
        return max1

    dfs(0, -1)
    return diameter


edges = [(0, 1, 3), (1, 2, 4), (1, 3, 2), (3, 4, 7)]
print(weighted_tree_diameter(5, edges))  # 3 + 2 + 7 = 12? Path 0-1-3-4
# Actually: 0->1 (3) + 1->3 (2) + 3->4 (7) = 12, or 2->1 (4) + 1->3 (2) + 3->4 (7) = 13
# Output: 13

---

6. Counting Paths / Subtree Sizes

A building block for many tree DP problems is computing subtree sizes. This is the simplest tree DP.

def compute_subtree_sizes(n, edges, root=0):
    """Compute the size of every subtree."""
    adj = defaultdict(list)
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)

    size = [1] * n

    def dfs(node, parent):
        for child in adj[node]:
            if child == parent:
                continue
            dfs(child, node)
            size[node] += size[child]

    dfs(root, -1)
    return size


edges = [(0, 1), (0, 2), (1, 3), (1, 4)]
sizes = compute_subtree_sizes(5, edges)
print(sizes)  # [5, 3, 1, 1, 1]

Subtree sizes are useful for problems like “sum of distances in a tree” (LeetCode 834), which we solve next with rerooting.

---

7. The Rerooting Technique

Problem: given a tree, compute some DP value for every possible root. Naive approach: run tree DP n times, giving O(n^2). Rerooting brings it down to O(n).

Idea:

1. Root the tree at node 0 and compute the DP bottom-up (first DFS).
2. In a second DFS, “reroot” from parent to child. When we move the root from u to its child v, we remove v’s contribution from u and add u’s updated answer to v.

Example: Sum of Distances in a Tree (LC 834)

Given a tree of n nodes, return an array where answer[i] is the sum of distances from node i to all other nodes.

def sum_of_distances_in_tree(n, edges):
    """
    LeetCode 834: Sum of Distances in Tree.

    Returns an array ans where ans[i] = sum of dist(i, j) for all j.
    """
    adj = defaultdict(list)
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)

    size = [1] * n   # subtree size
    ans = [0] * n    # sum of distances

    # DFS 1: compute subtree sizes and ans[0]
    def dfs1(node, parent, depth):
        ans[0] += depth  # distance from root to this node
        for child in adj[node]:
            if child == parent:
                continue
            dfs1(child, node, depth + 1)
            size[node] += size[child]

    dfs1(0, -1, 0)

    # DFS 2: reroot — move root from parent to child
    def dfs2(node, parent):
        for child in adj[node]:
            if child == parent:
                continue
            # When we move root from node to child:
            # - size[child] nodes get 1 closer (they're in child's subtree)
            # - (n - size[child]) nodes get 1 farther
            ans[child] = ans[node] - size[child] + (n - size[child])
            dfs2(child, node)

    dfs2(0, -1)
    return ans


n = 6
edges = [(0, 1), (0, 2), (2, 3), (2, 4), (2, 5)]
print(sum_of_distances_in_tree(n, edges))
# Output: [8, 12, 6, 10, 10, 10]

Time: O(n) — two DFS passes. Space: O(n).

Why Rerooting Works

When you shift the root from node u to its child v:

• Every node in v’s subtree is now one edge closer to the root.
• Every node outside v’s subtree is one edge farther.
• So ans[v] = ans[u] - size[v] + (n - size[v]).

This generalises to many problems: sum of depths, number of nodes at even distance, maximum distance, etc.

---

8. General Rerooting Template

Here is a more general rerooting template that works for problems beyond sum of distances.

def rerooting_template(n, edges, base, combine, finalize):
    """
    General rerooting DP template.

    Args:
        n: number of nodes
        edges: list of (u, v) pairs
        base: identity element for combine
        combine(accumulated, child_dp): merge a child's result
        finalize(accumulated, node): compute the final dp for a node

    Returns:
        dp_all[node] for each node as root
    """
    adj = defaultdict(list)
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)

    # Phase 1: root at 0, compute dp bottom-up
    dp = [base] * n
    order = []        # topological order
    parent = [-1] * n
    visited = [False] * n

    # BFS to get order
    from collections import deque
    queue = deque([0])
    visited[0] = True
    while queue:
        u = queue.popleft()
        order.append(u)
        for v in adj[u]:
            if not visited[v]:
                visited[v] = True
                parent[v] = u
                queue.append(v)

    # Process in reverse BFS order (leaves first)
    for u in reversed(order):
        acc = base
        for v in adj[u]:
            if v == parent[u]:
                continue
            acc = combine(acc, dp[v])
        dp[u] = finalize(acc, u)

    # Phase 2: reroot
    dp_all = dp[:]

    for u in order:
        # Compute prefix and suffix of children contributions
        children = [v for v in adj[u] if v != parent[u]]
        k = len(children)
        prefix = [base] * (k + 1)
        suffix = [base] * (k + 1)

        for i in range(k):
            prefix[i + 1] = combine(prefix[i], dp[children[i]])
        for i in range(k - 1, -1, -1):
            suffix[i] = combine(suffix[i + 1], dp[children[i]])

        for i, v in enumerate(children):
            # Remove v's contribution, add parent's contribution
            without_v = combine(prefix[i], suffix[i + 1])
            if parent[u] != -1:
                without_v = combine(without_v, dp_all[u])
            dp_rerooted = finalize(without_v, u)
            # Now v becomes root: its dp includes the "upward" subtree
            dp_all[v] = finalize(combine(dp[v], dp_rerooted), v)

    # Note: this is a simplified template. The exact combine/finalize
    # logic depends on the problem.
    return dp_all

---

9. DP on Trees — Common Patterns Summary

Pattern	States per Node	Example Problems
Include/exclude	2	Max independent set, House Robber III
Depth/height	1	Tree diameter, farthest node
Subtree sum	1	Subtree sizes, sum of values
Rerooting	2 passes	Sum of distances, min height trees
Path counting	varies	Count paths with sum K

---

10. Counting Nodes at Distance K

Another common tree DP problem: count how many pairs of nodes are at distance exactly K.

def count_pairs_at_distance_k(n, edges, k):
    """
    Count unordered pairs of nodes at distance exactly k.
    Uses centroid decomposition concept simplified for tree DP.
    """
    adj = defaultdict(list)
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)

    total_pairs = 0

    def dfs(node, parent):
        nonlocal total_pairs
        # cnt[d] = number of nodes at depth d in this subtree
        cnt = [0] * (k + 1)
        cnt[0] = 1  # the node itself is at depth 0

        for child in adj[node]:
            if child == parent:
                continue
            child_cnt = dfs(child, node)

            # Count pairs: one from existing subtrees, one from child's subtree
            for d in range(k):
                if k - 1 - d >= 0 and k - 1 - d <= k:
                    # depth d in existing + depth (k-1-d) in child's subtree
                    # +1 for the edge to child
                    need = k - 1 - d
                    if need < len(child_cnt):
                        total_pairs += cnt[d] * child_cnt[need]

            # Merge child's counts (shifted by 1 for the edge)
            for d in range(k):
                if d + 1 <= k:
                    cnt[d + 1] += child_cnt[d]

        return cnt

    dfs(0, -1)
    return total_pairs


edges = [(0, 1), (1, 2), (2, 3), (1, 4)]
print(count_pairs_at_distance_k(5, edges, 2))
# Pairs at distance 2: (0,2), (0,4), (2,4), (3,1) -> should be 4

---

11. Iterative Tree DP (Avoiding Stack Overflow)

For large trees (n > 10^5), Python’s default recursion limit is a problem. Here is how to convert tree DP to iterative using an explicit stack.

def max_independent_set_iterative(n, edges, weights):
    """Iterative version of max independent set on a tree."""
    adj = defaultdict(list)
    for u, v in edges:
        adj[u].append(v)
        adj[v].append(u)

    # BFS to find processing order and parents
    from collections import deque
    parent = [-1] * n
    order = []
    visited = [False] * n
    queue = deque([0])
    visited[0] = True

    while queue:
        u = queue.popleft()
        order.append(u)
        for v in adj[u]:
            if not visited[v]:
                visited[v] = True
                parent[v] = u
                queue.append(v)

    # Process in reverse BFS order (leaves first)
    dp_exc = [0] * n  # exclude node
    dp_inc = [0] * n  # include node

    for u in reversed(order):
        dp_inc[u] = weights[u]
        for v in adj[u]:
            if v == parent[u]:
                continue
            dp_inc[u] += dp_exc[v]
            dp_exc[u] += max(dp_exc[v], dp_inc[v])

    return max(dp_exc[0], dp_inc[0])


n = 7
edges = [(0, 1), (0, 2), (1, 3), (1, 4), (2, 5), (2, 6)]
weights = [3, 4, 5, 3, 2, 1, 2]
print(max_independent_set_iterative(n, edges, weights))

Key trick: BFS gives you a topological order. Processing in reverse guarantees children are handled before parents — the same guarantee post-order gives recursively.

---

12. When to Use Tree DP

Use tree DP when:

• The problem is on a tree (or can be modelled as one).
• The answer at a node depends on answers at its children.
• You need to compute something for every node (rerooting).
• The problem involves paths in a tree (diameter, distances).

Do not use tree DP when:

• The graph has cycles — use general graph DP or BFS/DFS.
• A simpler BFS/DFS gives the answer without needing DP states.

---

13. Practice Problems

Problem	Platform	Key Technique
House Robber III (LC 337)	LeetCode	Include/exclude DP
Sum of Distances in Tree (LC 834)	LeetCode	Rerooting
Binary Tree Maximum Path Sum (LC 124)	LeetCode	Max path through node
Diameter of Binary Tree (LC 543)	LeetCode	Two deepest paths
Tree Painting (CF 1187E)	Codeforces	Rerooting
Tree DP (CSES)	CSES	Include/exclude
Distance Sum (AtCoder)	AtCoder	Rerooting
Minimum Height Trees (LC 310)	LeetCode	Rerooting or leaf pruning

---

Big-O Summary

Algorithm	Time	Space
Max independent set	O(n)	O(n)
Tree diameter	O(n)	O(n)
Rerooting (2-pass)	O(n)	O(n)
Subtree sizes	O(n)	O(n)
Count pairs at distance K	O(nK)	O(nK)

Tree DP is one of the most satisfying patterns in competitive programming. The recursive structure of trees makes the DP transitions clean, and the rerooting technique turns what looks like an O(n^2) problem into O(n). Master these patterns and a large class of tree problems becomes approachable.