Dynamic Programming on Trees

Why trees are easy for DP

A tree has no cycles, so every node’s subtree is independent of everything outside it. That means a single depth-first traversal computes a value for each subtree, using only values already computed for its children. There is no ordering problem to solve — the recursion supplies the order.

The general shape:

static void dfs(int u, int parent) {
    // initialise dp[u] for a node with no children
    for (int v : adj.get(u)) {
        if (v == parent) continue;      // do not walk back up
        dfs(v, u);
        // combine dp[v] into dp[u]
    }
}

The if (v == parent) continue; line is what keeps an undirected tree from being traversed backwards. Without it the recursion immediately revisits the node it came from.

Subtree sizes and sums

The simplest useful case.

static int[] size;
static long[] sum;

static void dfs(int u, int parent) {
    size[u] = 1;
    sum[u] = value[u];
    for (int v : adj.get(u)) {
        if (v == parent) continue;
        dfs(v, u);
        size[u] += size[v];
        sum[u] += sum[v];
    }
}

size[u] counts the nodes in u’s subtree including itself; sum[u] totals their values. Both are O(N) overall since each edge is traversed once in each direction.

Choosing nodes with a constraint

Select a set of nodes with maximum total value such that no two selected nodes are adjacent.

Each node has two states: selected or not. If a node is selected, none of its children may be; if it is not, each child is free to choose whichever is better.

static long[][] dp;                     // dp[u][0] = u not taken, dp[u][1] = u taken

static void dfs(int u, int parent) {
    dp[u][0] = 0;
    dp[u][1] = value[u];

    for (int v : adj.get(u)) {
        if (v == parent) continue;
        dfs(v, u);
        dp[u][0] += Math.max(dp[v][0], dp[v][1]);   // child free to choose
        dp[u][1] += dp[v][0];                        // child must be skipped
    }
}

// answer: Math.max(dp[root][0], dp[root][1])

This is the template for most tree DP. The state is “this node, plus whatever small amount of information the constraint needs about it”, and the transition combines children independently.

Tree diameter

The longest path between any two nodes. For each node, the two deepest downward paths through it combine into a candidate.

static int best = 0;

static int depth(int u, int parent) {
    int deepest = 0, second = 0;
    for (int v : adj.get(u)) {
        if (v == parent) continue;
        int d = depth(v, u) + 1;
        if (d > deepest) { second = deepest; deepest = d; }
        else if (d > second) { second = d; }
    }
    best = Math.max(best, deepest + second);
    return deepest;
}

Each node returns only its deepest downward reach, but internally considers the top two, because a path may turn at that node rather than continue upward. Tracking only the deepest would miss every path that bends.

There is also a two-BFS method: the farthest node from any start is an endpoint of some diameter, and the farthest node from that is the other endpoint. Both are O(N); the DP version generalizes better.

Rerooting

Sometimes the question must be answered for every node as the root — for example, “for each node, the sum of distances to all other nodes”. Running a fresh DFS from each node is O(N²), too slow for large N.

Rerooting computes it in O(N) with two passes. The first collects information from below; the second pushes down what each node’s parent contributes from above.

static long[] below;                    // sum of distances within u's subtree
static long[] answer;                   // sum of distances to all nodes
static int[] size;
static int n;

static void down(int u, int parent) {
    size[u] = 1;
    below[u] = 0;
    for (int v : adj.get(u)) {
        if (v == parent) continue;
        down(v, u);
        size[u] += size[v];
        below[u] += below[v] + size[v];   // each subtree node is one edge further
    }
}

static void up(int u, int parent) {
    for (int v : adj.get(u)) {
        if (v == parent) continue;
        // moving the root from u to v: size[v] nodes get closer, the rest farther
        answer[v] = answer[u] - size[v] + (n - size[v]);
        up(v, u);
    }
}

// answer[root] = below[root]; then up(root, 0)

answer[v] = answer[u] - size[v] + (n - size[v]);

Recursion depth

A tree with 200,000 nodes shaped like a path gives a recursion 200,000 deep, which overflows Java’s default stack. Either rewrite the traversal iteratively, or run the solution on a thread with a larger stack:

public static void main(String[] args) {
    new Thread(null, Main::solve, "main", 1 << 26).start();
}

This is the single most common reason a correct tree DP fails on the large test cases. See Debugging Under Time Pressure.

Related