Sum of Distances in Tree
Understanding the Problem
We have an unweighted tree with n nodes, and we need to compute, for every node, the total distance to all other nodes. The direct solution runs a BFS from each node to find distances to every other node, but with n up to 30,000 that is O(n^2), around 900 million operations.
The better idea is to compute the answer for all nodes without starting from scratch each time. If we already know the answer for one node, we can derive the answer for its neighbor in O(1). That technique, computing the answer at one root and then shifting it to adjacent nodes, is called rerooting, and it brings the total down to O(n).
Key Constraints:
1 <= n <= 3 * 10^4→ O(n^2) is about 900 million operations, too slow for strict time limits. We need an O(n) solution.edges.length == n - 1→ This is a tree: connected, no cycles. Every pair of nodes has exactly one path between them, so distance is well defined.- The tree is unweighted, so distance equals the number of edges on the path.
- The largest possible answer occurs in a path graph at an endpoint: 0 + 1 + ... + (n-1) = n(n-1)/2, about 4.5 10^8 for n = 30,000. That fits comfortably in a signed 32-bit integer (max about 2.1 10^9), so plain
intis safe in every language here.
Approach 1: BFS from Every Node
Intuition
For each node, run a BFS to compute the distance to every other node, then sum those distances. Repeat for all n nodes.
For a single source, BFS visits every other node exactly once and assigns each its distance from the source, so one BFS gives that node's total in O(n). Running it from all n sources is O(n^2). This is a useful baseline because it directly encodes the definition of the answer, and it establishes the correct results we will reproduce faster in the next approach.
Algorithm
- Build an adjacency list from the edges.
- For each node
ifrom 0 to n-1: - Run BFS starting from node
i. - Sum up the distances to all other nodes.
- Store the sum in
answer[i]. - Return the answer array.
Example Walkthrough
Code
- Time Complexity: O(n^2). We run BFS n times, and each BFS visits all n nodes. Total work is O(n * n) = O(n^2).
- Space Complexity: O(n). The adjacency list takes O(n) space (since there are n-1 edges). Each BFS uses O(n) for the dist, visited arrays, and queue.
This repeats a full BFS from every node and discards the work each run does. The distance sums for neighboring nodes are closely related, and the next approach exploits that relationship: compute the answer once for a chosen root, then derive each neighbor's answer from it in constant time.
Approach 2: Rerooting DP (Optimal)
Intuition
Root the tree at node 0 and suppose we know answer[p] for a node p. Take one of p's children c. Moving the reference point from p to c shifts every distance by exactly one edge, because p and c are adjacent. Split the nodes by which side of the edge p-c they lie on:
- The nodes in
c's subtree are reached throughc, so moving tocbrings each of them 1 closer. - Every other node is reached through
p, so moving tocputs each of them 1 farther.
If count[c] is the number of nodes in c's subtree, then count[c] distances drop by 1 and n - count[c] distances rise by 1:
answer[c] = answer[p] - count[c] + (n - count[c]) = answer[p] + n - 2 * count[c]
This needs two pieces of information: the answer for the root and the subtree size of every node. A first DFS computes both. A second DFS walks down from the root applying the formula, giving each node's answer in O(1). The whole algorithm runs in O(n).
The formula uses count[c], the subtree size measured when the tree was rooted at node 0 in DFS 1. DFS 2 never recomputes it. This is safe because count[c] counts the nodes on c's side of the edge p-c, and removing that edge always splits the tree into the same two components no matter which endpoint we call the root. So count[c] is exactly the number of nodes 1 closer when we move from p to c, and n - count[c] is the number 1 farther. Both quantities depend only on the edge, not on the rooting, so the value carried over from DFS 1 remains correct at every step.
Algorithm
- Build an adjacency list from the edges.
- Establish a parent pointer and a visit order. Run one iterative DFS from node 0, pushing unvisited neighbors and recording
parent[node]and the order nodes are popped. Because a child is always popped after its parent, the reverse of this order is a valid bottom-up (post-order) processing sequence. - DFS 1 (bottom-up): Walk the order in reverse. For each node, set
count[node] = 1 + sum(count[child])over its children, and accumulateanswer[0]asanswer[node] += answer[child] + count[child]. After this pass,answer[0]holds the sum of distances from the root, andcount[i]holds every subtree size. - DFS 2 (top-down): Walk the order from index 1 onward (skipping the root). For each node
cwith parentp, applyanswer[c] = answer[p] + n - 2 * count[c]. The parent is always processed before the child, soanswer[p]is final whencis computed. - Return the answer array.
The two passes use an explicit stack and array iteration rather than recursion. With n up to 30,000, a path-shaped tree has a recursion depth of 30,000, which overflows the default call stack in several languages. The iterative form keeps the depth bounded by the heap-allocated stack and avoids that failure.
Example Walkthrough
Code
- Time Complexity: O(n). Building the adjacency list takes O(n) (there are n-1 edges). DFS 1 visits each node once: O(n). DFS 2 visits each node once: O(n). Total: O(n).
- Space Complexity: O(n). The adjacency list, count array, answer array, parent array, order array, and visited array each take O(n) space. The explicit DFS stack also holds at most n nodes. Because the traversal is iterative, there is no O(h) recursion stack to overflow on a skewed tree.