Skip to content
CCC Python Course

Tree dynamic programming

Module
M6.5
Lesson
1 of 1
Reading time
5 min

In this lesson

  • Recognize problems solvable by computing values on trees with subproblem overlap.
  • Root a tree and define a subproblem on each subtree.
  • Compute solutions bottom-up from leaves to root.
  • Rerooting: compute solutions at every node without recomputing from scratch.

Trees are simpler than general graphs but complex enough to need dynamic programming. A tree has no cycles. Each node except the root has exactly one parent. This structure means a subtree is independent of everything outside it. Compute the answer on each subtree, combine answers from children, and propagate upward.

Dynamic programming on trees solves problems about selecting subsets of nodes, counting paths, assigning values, or optimizing over all subtrees. The key is breaking the problem into independent subproblems on each subtree.

Rooting the tree

A tree has no inherent root, but for DP you pick one. Say the root is node 0. Every other node now has one parent and zero or more children. A subtree at node u contains u and all its descendants.

Define a subproblem: dp[u] is the answer to the problem restricted to the subtree rooted at u. To compute dp[u], use dp[v] for all children v of u, plus information about u itself.

Many problems ask for a single value: the maximum, minimum, or count. Some ask for multiple values: dp[u][0] and dp[u][1] for "u is not selected" and "u is selected".

An example: maximum independent set

Select a subset of nodes such that no two are adjacent, and count how many you select. In a tree, this means if you select a node, you cannot select any of its children or parent.

Define dp[u][0] as the maximum size when u is not selected, and dp[u][1] as the maximum size when u is selected.

If u is not selected, every child v can be selected or not. Take the maximum: dp[u][0] = sum(max(dp[v][0], dp[v][1])).

If u is selected, no child can be selected. So dp[u][1] = 1 + sum(dp[v][0]).

examples/tree_dp.py
import sys

def dfs(node, parent, adj, dp):    dp[node][0] = 0    dp[node][1] = 1
    for child in adj[node]:        if child != parent:            dfs(child, node, adj, dp)            dp[node][0] += max(dp[child][0], dp[child][1])            dp[node][1] += dp[child][0]

def main() -> None:    input_data = sys.stdin.read().split()    if not input_data:        return
    idx = 0    n = int(input_data[idx])    idx += 1
    adj = [[] for _ in range(n)]    for _ in range(n - 1):        u = int(input_data[idx])        v = int(input_data[idx + 1])        idx += 2        adj[u].append(v)        adj[v].append(u)
    dp = [[0, 0] for _ in range(n)]    dfs(0, -1, adj, dp)
    result = max(dp[0][0], dp[0][1])    sys.stdout.write(str(result) + "\n")

if __name__ == "__main__":    main()

Input

6
0 1
0 2
1 3
1 4
2 5

Output

4
Maximum independent set on a tree

The program reads a tree and computes the maximum independent set size using DP.

Why tree DP works efficiently

A tree with n nodes has n - 1 edges and no cycles. This means each node has a unique path to any other node. A subtree is independent: removing any node splits the tree into disconnected components, and the DP answer in each component depends only on what is inside that component.

Because each node appears in the subtree of only one child of its parent, when you combine answers upward, you compute each subproblem once. The total work is O(n) for the DP computation (one per node), plus O(n) for reading and output. This is much faster than trying all 2^n subsets.

A second example: counting paths with a target sum

Count the number of paths in a tree that sum to a target value k. A path is a sequence of nodes along the tree. We want paths in any direction and of any length.

Define dp[u] as a dictionary mapping "sum from u downward" to the count of paths. For a leaf, dp[u] = {value[u]: 1}. For an internal node u, process its children one at a time. Before merging in a new child's dictionary, check every pair of an already-merged sum s1 and a new sum s2 from that child: if s1 + s2 == k, add the count to the result. Then merge the new child's sums, each shifted by value[u], into dp[u].

This approach counts paths that pass through u, combining any two of its children (not just two, if u has more), and paths that end at descendants. By using dictionaries to track only reachable sums, the algorithm avoids exponential blowup.

Rerooting

Sometimes you need the answer not just at one root, but at every node. Rerooting computes this efficiently without repeating all computation.

First, compute DP values with the tree rooted at node 0. Then, reroot at each neighbor: use the DP values to compute what the answer would be if that neighbor were the root. Combine downward information from the original root with upward information from the children.

The classic rerooting technique computes two passes: once up the tree, once down. The first pass computes answers on each subtree. The second pass applies information from parent to child, and the answer at each node incorporates both directions.

Common mistakes

One mistake is to forget that you must process nodes in the right order. For bottom-up DP on trees, process children before parents. Use post-order DFS: visit all descendants and compute their DP values before processing the node itself. If you process a node before its children, the children's DP values are not yet ready, and your result is wrong.

Another mistake is to overcomplicate the state space. Some problems need dp[u][k] for each possible k selected from the subtree. But if k can range from 0 to n, you have n² states and O(n³) time. Before accepting this, check whether you can use a dictionary to store only reachable states, or whether a simpler formulation like "selected or not selected" suffices.

A third mistake is to forget the contribution of the node itself. When computing dp[u], include value[u] or some property of u. Many learners compute only the contributions from children and forget the node's own data, leading to off-by-one or missing-data errors.

A fourth mistake is to confuse rerooting with simply recomputing. Rerooting requires careful bookkeeping: when you move the root, answers for old ancestors become descendants, and their contribution changes. Recompute explicitly by undoing the connection to the old parent and adding the connection to the new root. Naive attempts to adjust DP values on the fly often miss dependencies and produce incorrect answers.

Recap

Tree dynamic programming solves problems by breaking them into independent subproblems on each subtree. Root the tree, define a subproblem on each subtree, and combine answers from children upward. The simple bottom-up approach is the foundation. Rerooting extends this to compute answers for every node as root without redoing all computation. Tree DP is useful and relatively simple compared to general DP on DAGs or arbitrary graphs.

Practice

Try this on the judge. The link opens the problem on WMOJ.

  1. 2022 S5
    Good Influencers (opens on WMOJ in a new tab) WMOJ

    Keep several states per node and combine children in post-order to satisfy a tree-wide condition.