834 · Sum of Distances in Tree

The one problem that proves a single DFS is not enough. Pass 1 computes subtree sizes and the root's answer bottom-up; pass 2 derives every child's answer from its parent's in O(1): ans[c] = ans[p] + n − 2 · size[c]. Derive that formula — do not memorise it.

1 — The problem

An undirected tree with n nodes. For every node, return the sum of distances from it to all others. Doing a BFS from each node is O(n²) and times out at n = 3 × 10⁴.

Everything before this in Pattern 2 has computed one answer, at the root. Here every node needs its own, and that changes the shape of the solution: information has to flow down as well as up, and a single postorder cannot do both.

This shape — measure once bottom-up, then reroot top-down — is sub-variant L, and it is the answer to any problem phrased as "compute X for every node" where X depends on the whole tree rather than just a subtree.

2 — Deriving the rerooting formula

This is the part to be able to reproduce. Move the root from p to its child c, and ask what happens to each node's distance:

• Every node inside c's subtree gets one step closer. There are size[c] of them, so the total falls by size[c].
• Every node outside it gets one step further. There are n − size[c] of them, so the total rises by that much.

Therefore ans[c] = ans[p] − size[c] + (n − size[c]), which tidies to:

ans[c] = ans[p] + n − 2 · size[c]

Two sanity checks worth doing in your head. A leaf has size = 1, so ans[leaf] = ans[parent] + n − 2 — one node got closer, n − 1 got further, net +n − 2. And a child whose subtree holds exactly half the tree gives ans[c] = ans[p], which is right: the gains and losses cancel.

After pass 1, only ans[0] = 8 is a real answer — every other entry holds the sum of distances within its own subtree, which is not what was asked. Pass 2 overwrites them all, parent before child, so each read of ans[u] is already correct.

2.1 Why line 6 accumulates ans[v] + size[v]

In pass 1, ans[u] means "sum of distances from u to everything in u's subtree". Adding child v's subtree contributes ans[v] — the distances measured from v — plus one extra step for each of the size[v] nodes, to get from u down to v. Hence + size[v].

The quantity therefore changes meaning between the passes: subtree-local after pass 1, whole-tree after pass 2. Reusing the same array is conventional and efficient, and it is also the main source of confusion when reading this code — the entry means two different things depending on when you look at it.

2.2 Practical notes

• The tree is given as an edge list, undirected. Build an adjacency list; the v != p guard is what stops the recursion walking back up. There is no left/right and no natural root — pick 0.
• Recursion depth can reach n on a path graph. At 3 × 10⁴ that risks a stack overflow in Java; an explicit stack or an iterative order is the safe version if the constraints grow.
• Pass 2 must be preorder. ans[v] is computed from ans[u], so the parent has to be final first. Running it postorder reads uninitialised values and fails silently.

3 — Complexity and edge cases

• Time O(n) — two passes. Space O(n) for the adjacency list and the two arrays, plus O(h) recursion.
• n = 1: the answer is [0]. Both passes do nothing.
• A star (one centre, n − 1 leaves): the centre scores n − 1, every leaf scores 2n − 3. Good formula check.
• A path: the answers form a parabola with its minimum at the middle — which is what LC 310, Minimum Height Trees, is really finding.
• Common bug: forgetting + size[v] in pass 1. The root's answer comes out too small and every derived answer inherits the error.
• Common bug: using 2 * size[u] instead of 2 * size[v] in the reroot. It is the child's subtree that moves closer.
• Common bug: running pass 2 before pass 1 has finished, or merging them into one traversal. They cannot be merged — that is the entire point of the problem.

4 — Reference implementation

private List<Integer>[] g;
private int[] size, ans;
private int n;

public int[] sumOfDistancesInTree(int n, int[][] edges) {
    this.n = n;
    g = new List[n];
    for (int i = 0; i < n; i++) g[i] = new ArrayList<>();
    for (int[] e : edges) { g[e[0]].add(e[1]); g[e[1]].add(e[0]); }
    size = new int[n];
    ans  = new int[n];
    post(0, -1);
    pre(0, -1);
    return ans;
}

private void post(int u, int p) {
    size[u] = 1;
    for (int v : g[u]) if (v != p) {
        post(v, u);
        size[u] += size[v];
        ans[u]  += ans[v] + size[v];   // +1 step for each node in v's subtree
    }
}

private void pre(int u, int p) {
    for (int v : g[u]) if (v != p) {
        ans[v] = ans[u] + n - 2 * size[v];   // reroot from u to v
        pre(v, u);                          // parent must be final first
    }
}

Related: LC 310, Minimum Height Trees answers a similar "every node needs a number" question by peeling leaves topologically instead of rerooting — cheaper, and a good contrast for when the full set of values is not actually required.