Tree DP and rerooting: answers for every root in linear time

Why this matters when you build with AI

Ask an AI tool for "the sum of distances from every vertex of a tree" and you will often get a correct recursive function that is too slow. It runs one DFS per root. On a path of 200000 vertices that is 200000 traversals of 200000 vertices each, about 4 * 10^10 steps. The answers on your small samples look right, so the mistake shows up only as a time limit error on the judge.

A second common failure is recursion. A DFS written with recursion on a path of 200000 vertices goes 200000 calls deep. Whether that crashes depends on the stack size of the machine and on the size of each call frame, so the same code can pass on one machine and fail on another.

You can catch both. Count the work from the limits before you read the code. Then write a brute force that is slow but obviously right, and compare the two on many small random trees. This lesson builds that habit with the two standard tools for "answer for every root": subtree DP and rerooting.

A prompt that states the limits and asks for the right shape of answer:

Write C++17 code that, for every vertex r of a tree with up to 200000 vertices,
computes the sum of distances from r to all other vertices. Use an O(n) method,
not one BFS per root. Do not use recursion, because the tree can be a path of
200000 vertices. Use long long for the sums. Explain the two passes you use and
why each one visits the vertices in the order it does.

A good answer has two passes over one BFS order, one pass from the leaves up and one from the root down, and uses long long for the sums. A bad answer runs a DFS inside a loop over the roots, or recurses on the tree. Check each point by reading the code, then run it against a brute force (section "Check them against each other" below).

Terms used

A tree has n vertices and n - 1 edges, and exactly one path between any two vertices. A tree has no cycles.
Rooting a tree at a vertex means picking it as the top. Every other vertex then has one parent (the next vertex on the way to the root) and any number of children.
The subtree of v is v and every vertex below it. Its size is the number of vertices in it.
The distance between two vertices is the number of edges on the path between them.
The eccentricity of r is the largest distance from r to any vertex. The smallest eccentricity over all r is the radius of the tree.
DP means dynamic programming here: a value for each vertex is computed from the values of its children, or of its neighbours, so each value is computed once.
Rerooting means moving the root from one vertex to a neighbour and updating the answer in constant time, instead of starting over.
A prefix combine of a list is the result for the first i items. A suffix combine is the result for the last items. Taking the prefix before i and the suffix after i gives the result for everything except item i, without dividing or removing anything.
A BFS order is the list of vertices when you visit them in breadth-first order from the root. In this order every parent comes before its children.

The idea: one tree, two passes

All code in this lesson uses the same small tree. It has 7 vertices, numbered 0 to 6:

0 - 1      1 - 3      1 - 4
0 - 2      2 - 5      5 - 6

Draw it as a picture and it looks like this. Vertex 0 has two children, 1 and 2. Vertex 1 has two children, 3 and 4. Vertex 2 has one child, 5, and 5 has one child, 6.

The answer for vertex 0 is the sum of its distances to all other vertices: 1 + 1 + 2 + 2 + 2 + 3 = 11. Moving the root from 0 to its child 1 changes each distance by exactly 1, so the answer for vertex 1 is the answer for vertex 0 plus a correction. This is the whole idea of rerooting.

The two passes are:

Subtree pass, from the leaves up. For each vertex compute its subtree size and the sum of distances from it to the vertices in its subtree. Each vertex needs only its children's values.
Rerooting pass, from the root down. Start with the answer for the root. For each child, compute its answer from its parent's answer.

Each pass visits every vertex once, so the total work is O(n).

Subtree DP: answers that flow up

Let sz[v] be the size of the subtree of v, and down[v] be the sum of the distances from v to the vertices in its subtree.

For a child c of v, every vertex in the subtree of c is one step further from v than from c. The subtree of c has sz[c] vertices, so the sum of their distances from v is down[c] + sz[c]. Adding this over all children gives down[v]:

sz[v]   = 1 + sum of sz[c] over children c
down[v] = sum of (down[c] + sz[c]) over children c

Vertex 3 is a leaf, so sz[3] = 1 and down[3] = 0. Vertex 6 is also a leaf. Vertex 5 has one child, 6, so sz[5] = 2 and down[5] = (0 + 1) = 1, the distance from 5 to 6. Vertex 1 has children 3 and 4, so sz[1] = 3 and down[1] = (0+1) + (0+1) = 2.

The rule needs the children to be finished before the parent. Any order where every child comes before its parent works. The reverse of a BFS order is one.

Rerooting: answers that flow down

Once down is known for the root, the answer for the root is down[0]. Now look at one edge, from a parent p to a child v, and move the root across it.

The sz[v] vertices in the subtree of v get one step closer to the root. Their distances drop by 1 each.
The other n - sz[v] vertices get one step further. Their distances rise by 1 each.

So the answer for v is the answer for p, minus sz[v], plus n - sz[v]:

ans[v] = ans[p] - sz[v] + (n - sz[v])

For vertex 1 in the example, ans[0] = 11, sz[1] = 3, and n = 7. So ans[1] = 11 - 3 + 4 = 12. Count by hand: from vertex 1, the distances to 0, 2, 3, 4, 5 and 6 are 1, 2, 1, 1, 3 and 4. They add up to 12, which is the formula's value.

This is why the rerooting pass needs the subtree sizes from the first pass: the correction depends on sz[v].

The code for the sum

The two passes are in distance_sums below. The rerooting loop goes in BFS order, so the parent's answer is always ready before the child's.

Eccentricity: a maximum with prefix and suffix combines

A sum is easy because a correction can be added. A maximum is harder, because you cannot subtract a child's contribution from a maximum. The fix is to combine everything except one child, directly.

Let down_h[v] be the longest path from v going down into its subtree. It is 0 for a leaf. For each child c, the child offers the value down_h[c] + 1 to its parent.

Let up[v] be the longest path from v that first goes to its parent. Then the eccentricity of v is max(down_h[v], up[v]).

To find up[c] for a child c of p, you need the best of:

up[p], the way up from p,
the best child value of p, not counting c.

That last part is the important one. Put the values of p's children in a list. For child j, the best value excluding j is the maximum of two things: the maximum over the children before j (a prefix combine), and the maximum over the children after j (a suffix combine). Build both arrays once per vertex. Then each child gets its answer in constant time.

children of 0 offer: 1 + down_h of child 1, 1 + down_h of child 2

Using 0 as the starting value of every prefix and suffix means "stay at p", which is a path of length 0. The root has up = 0 for the same reason.

Iterative traversal order

A recursive DFS is the natural way to write this, and it causes the stack problem above. The code in this lesson never recurses. It does two things:

Build a BFS order from the root with a queue (a plain vector and an index is enough). Store the parent of each vertex.
Walk the order backwards for the subtree pass, so that each child is done before its parent. Walk it forwards for the rerooting pass, so that each parent is done before its child.

A queue works because each vertex is added once, when its parent is taken out. The order is not a DFS order, but it does not need to be: the two passes only need "children before parent" or "parent before children".

The C++17 code

The code is in two files. reroot.hpp holds the tree type, the BFS order, the sum and the eccentricity. main.cpp reads a tree and prints one line per vertex.

#pragma once
#include <algorithm>
#include <cstddef>
#include <vector>
 
// Tree stored as adjacency lists. Vertices are 0 .. n-1.
struct Tree {
    int n;
    std::vector<std::vector<int>> adj;
    explicit Tree(int n) : n(n), adj(n) {}
    void add_edge(int a, int b) {
        adj[a].push_back(b);
        adj[b].push_back(a);
    }
};
 
// Root the tree at 0 and return a BFS order: parents come before children.
// parent[0] = -1.
struct Rooted {
    std::vector<int> order;
    std::vector<int> parent;
};
 
inline Rooted root_at_zero(const Tree& t) {
    Rooted r;
    r.parent.assign(t.n, -1);
    std::vector<char> seen(t.n, 0);
    r.order.reserve(t.n);
    r.order.push_back(0);
    seen[0] = 1;
    for (std::size_t i = 0; i < r.order.size(); ++i) {
        int v = r.order[i];
        for (int w : t.adj[v]) {
            if (!seen[w]) {
                seen[w] = 1;
                r.parent[w] = v;
                r.order.push_back(w);
            }
        }
    }
    return r;
}
 
// sum[r] = sum over all v of dist(r, v), for every root r. Linear time.
inline std::vector<long long> distance_sums(const Tree& t) {
    const int n = t.n;
    Rooted R = root_at_zero(t);
    std::vector<long long> sz(n, 1);   // size of subtree of v
    std::vector<long long> down(n, 0); // sum of dist(v, u) for u in subtree of v
    // Subtree pass: children are handled before their parent (reverse BFS order).
    for (int i = n - 1; i >= 1; --i) {
        int v = R.order[i];
        int p = R.parent[v];
        sz[p] += sz[v];
        down[p] += down[v] + sz[v]; // every vertex under v is one step further from p
    }
    // Rerooting pass: moving the root from p to child v.
    // The sz[v] vertices under v get 1 closer, the other n - sz[v] get 1 farther.
    std::vector<long long> ans(n, 0);
    ans[0] = down[0];
    for (int i = 1; i < n; ++i) {
        int v = R.order[i];
        int p = R.parent[v];
        ans[v] = ans[p] - sz[v] + (n - sz[v]);
    }
    return ans;
}
 
// ecc[r] = largest dist(r, v) over all v (the eccentricity of r), for every root r.
inline std::vector<int> eccentricities(const Tree& t) {
    const int n = t.n;
    Rooted R = root_at_zero(t);
    std::vector<int> down(n, 0); // longest path from v going down into its subtree
    for (int i = n - 1; i >= 1; --i) {
        int v = R.order[i];
        int p = R.parent[v];
        down[p] = std::max(down[p], down[v] + 1);
    }
    // up[v] = longest path from v that starts with the edge v -> parent(v).
    // up[0] = 0 means "stay at the root", which is a path of length 0.
    std::vector<int> up(n, 0);
    for (int i = 0; i < n; ++i) {
        int p = R.order[i];
        // Children of p, and the value each child offers upward: down[c] + 1.
        std::vector<int> kids;
        for (int c : t.adj[p]) {
            if (c != R.parent[p]) kids.push_back(c);
        }
        int k = static_cast<int>(kids.size());
        std::vector<int> pre(k + 1, 0), suf(k + 1, 0);
        for (int j = 0; j < k; ++j) pre[j + 1] = std::max(pre[j], down[kids[j]] + 1);
        for (int j = k - 1; j >= 0; --j) suf[j] = std::max(suf[j + 1], down[kids[j]] + 1);
        // For child j, the best way up leaves out child j: combine pre[j] and suf[j+1].
        for (int j = 0; j < k; ++j) {
            int c = kids[j];
            up[c] = 1 + std::max(up[p], std::max(pre[j], suf[j + 1]));
        }
    }
    std::vector<int> ecc(n);
    for (int v = 0; v < n; ++v) ecc[v] = std::max(down[v], up[v]);
    return ecc;
}

#include <cstdio>
#include <vector>
 
#include "reroot.hpp"
 
// Input: first line n. Then n - 1 lines "a b", one undirected edge each.
// Output: one line per vertex r: "r sumdist ecc".
int main() {
    int n;
    if (std::scanf("%d", &n) != 1) return 0;
    Tree t(n);
    for (int i = 0; i + 1 < n; ++i) {
        int a, b;
        if (std::scanf("%d %d", &a, &b) != 2) return 1;
        t.add_edge(a, b);
    }
    std::vector<long long> s = distance_sums(t);
    std::vector<int> e = eccentricities(t);
    for (int r = 0; r < n; ++r) std::printf("%d %lld %d\n", r, s[r], e[r]);
    return 0;
}

Input for the example tree, saved as example.txt:

Real output of the program on that input, compiled with g++ -O2 -std=c++17 -Wall -Wextra:

The first number is the vertex, the second is the sum of distances, the third is the eccentricity. You can check two of them by hand. For vertex 0 the sum is 11 and the farthest vertices are 3, 4 and 6 at distance 3, so the eccentricity is 3. For vertex 6 the distances are 1, 2, 3, 4, 5 and 5, so the sum is 20 and the eccentricity is 5.

A Python version

The same two passes fit in a short Python function. It is useful for reading the logic and for small inputs. For limits near 200000, use C++.

Python · editable

Real output, run with python3 -I (Python 3.14.4 on the machine used for this lesson):

[11, 12, 12, 17, 17, 15, 20]

These match the sums in the C++ output above.

Complexity

Time: O(n) for each of the two passes, plus O(n) for the BFS, so O(n) in total. The eccentricity code also builds a small list of children for each vertex, and the total size of those lists is n - 1, so it stays O(n).
Memory: O(n) for the adjacency lists, the order, the parents and the answers.
Recursion: none.

The naive method, one BFS per root, costs O(n^2). For n = 200000 that is 4 * 10^10 steps, which is far above what a judge allows.

Check them against each other

The file stress.cpp builds random trees and compares the two functions with a brute force. The brute force runs one BFS from each root and adds up the distances and the largest distance. It is O(n^2), which is fine for n up to 40.

The random trees come in three shapes: a random recursive tree (each new vertex joins a random earlier vertex), a path, and a tree where each new vertex joins one of the first two vertices. Every tree gets its labels shuffled, so the root 0 is not always a special vertex.

Real output, with two seeds:

seed 7: 3000 trees, 61991 (tree, root) pairs, all match
seed 99: 3000 trees, 61967 (tree, root) pairs, all match

Each line says how many trees and how many (tree, root) pairs were checked. Both the sum and the eccentricity had to match for every root.

The eccentricity values are checked too, so a change to one of the two passes shows up in the output.

Common mistakes

Mistake 1: one BFS per root. Symptom: samples pass and a large test runs past the time limit. The cause is O(n^2) work. The fix is the two-pass method above. Check the loop structure first: a loop over roots containing a traversal is the warning sign.

Mistake 2: recursion on a long path. Symptom: a crash (a runtime error from stack overflow) on a test with a long chain, while small random trees pass. The depth of recursion equals the length of the longest chain, and whether that overflows depends on the judge's stack size. Use the BFS order instead.

Mistake 3: wrong sign or wrong size in the rerooting step. Symptom: wrong answers for every vertex except the root, off by an amount that depends on the subtree size. The stress test shows it. I changed the rerooting line to ans[v] = ans[p] + n - sz[v];, which drops the - sz[v] term. The stress test stopped at the first small tree:

MISMATCH test 0 root 1: sum 108 vs 94, ecc 13 vs 13

Write the formula on paper for one edge, with the sizes, before you code it.

Mistake 4: reading a parent's answer before it is set. Symptom: a child gets a value computed from an unset starting value. If the loop visits a child before its parent, ans[p] has not been set yet. The BFS order prevents this when you walk it forwards.

Mistake 5: not leaving out the child when combining for the maximum. Symptom: the upward value of a child includes its own subtree, so the eccentricity is too large. I tested this by replacing the prefix and suffix combine with the maximum over all children. The stress test printed the first mismatch, and on the path 0 - 1 - 2 the output for vertex 1 was wrong:

MISMATCH test 0 root 1: sum 94 vs 94, ecc 15 vs 13

Correct output for that path (vertex, sum, eccentricity):

0 3 2
1 2 1
2 3 2

The mutant printed 1 2 3 for vertex 1, because it counted the path from vertex 1 down to itself as a way up.

Checklist

Write the sum of the limits first. If the naive method is n^2, use the passes.
Use a BFS order for both passes. Do not recurse on a tree that can be a long path.
For a sum over subtrees, the subtree pass goes children before parents.
For a rerooting sum, write the change for one edge with the sizes before you code it.
For a maximum over all roots, build prefix and suffix combines of the children's values, and leave out the child you are answering for.
Compare with a brute force on random trees, including paths and stars, and check every root.
Use long long for sums. A path of 200000 vertices can give sums near 2 * 10^10.

What was checked

All C++ code compiled with g++ 15.2.0 (Ubuntu 15.2.0-16ubuntu1) using -O2 -std=c++17 -Wall -Wextra, with no warnings after the final fix.
The stress test passed on 3000 random trees with seed 7 and 3000 with seed 99, comparing every root against the brute force. The counts above are from that run.
The two mutants shown in the mistakes section were run and the printed lines were copied from the output.
The hand-checked values (vertex 0 sum 11, vertex 1 sum 12, vertex 6 sum 20 and eccentricity 5) match the program output.
The Python sample was run with python3 -I.
I did not time the C++ code on a 200000-vertex tree. Running times depend on the machine.

Summary

Subtree DP computes a value for each vertex from its children, going up. Rerooting moves the answer from a parent to a child in constant time, going down. A sum can be corrected by adding and subtracting the subtree size. A maximum needs prefix and suffix combines, so a child can be left out. Build the order once, with a BFS, and both passes become loops that never recurse.

Practice

Problem 1: trace the example

Use the example tree from this lesson. Without running the code, find the eccentricity of vertex 6 and every vertex that is farthest from it. Then compare with the program output.

Show the answer

Hint: list the distance to each vertex by walking the path. Start with the neighbour 5.

Answer: the eccentricity of 6 is 5. The farthest vertices are 3 and 4, both at distance 5 (6, 5, 2, 0, 1, then 3 or 4). The sum of distances is 20, which matches the program output 6 20 5.

Problem 2: weighted edges

Each edge now has a positive weight w, and the distance is the sum of the weights on the path. Write the rerooting formula for the sum of distances, for a move from parent p to child v across an edge of weight w. Use sz[v] and n.

Show the answer

Hint: the vertices under v get closer by w, and the others get farther by w. Multiply the count by w.

Answer: ans[v] = ans[p] + w * (n - 2 * sz[v]). The subtree pass also needs down[p] += down[v] + w * sz[v]. The sum of weights can exceed 2^31 on a path, so keep long long.

Problem 3: radius from the eccentricities

The radius of a tree is the smallest eccentricity. Given the output of main on the example tree, which vertices are centres (the vertices with eccentricity equal to the radius), and what is the radius?

Show the answer

Hint: find the smallest third number in the output lines.

Answer: the radius is 3. The centres are vertices 0 and 2, both with eccentricity 3.

Problem 4: count the vertices at distance exactly k

Explain why the rerooting trick does not give the number of vertices at distance exactly k from every root in O(n) with the same two passes, and say what extra structure would be needed.

Show the answer

Hint: the sum changes by a fixed amount when the root moves across one edge. Ask whether the count of vertices at one distance changes by a fixed amount.

Answer: when the root moves across one edge, the distances change by 1 or -1 in a way that depends on how many vertices sit at each depth, not on one number like the subtree size. The two-pass method stores one number per vertex, so it cannot carry that information. Keeping the depth counts for every vertex can take O(n^2) memory in the worst case. This lesson does not give a method for this variant.

Problem 5: estimate the naive cost

A path of 200000 vertices is given. Count the steps of the one-BFS-per-root method, and say how many steps the two-pass method takes.

Show the answer

Hint: a BFS from one root visits every vertex, so one BFS costs n steps. There are n roots.

Answer: n roots times n steps is 200000 * 200000 = 4 * 10^10 steps for the naive method. The two-pass method does three linear sweeps (the BFS and the two passes), about 3 * 200000 = 6 * 10^5 steps. These are counts of steps, not measured times.