AlgoInsight

DSA problem

Sum of Distances in Tree

HardTree / Graph DPPattern: Tree / Graph DPtop300GoogleAmazonMetaMicrosoftBloomberg

There is an undirected connected tree with n nodes labeled from 0 to n - 1 and n - 1 edges. You are given the integer n and the array edges where edges[i] = [a_i, b_i] indicates that there is an edge between nodes a_i and b_i in the tree. Return an array answer of length n where answer[i] is the sum of the distances between node i and every other node in the tree.

A brute-force solution does a BFS or DFS from each node, costing O(n^2) time, which is too slow for n up to 3 * 10^4. The trick is the "re-rooting" technique: compute the answer at one anchor node in O(n), then derive each other node's answer in O(1) by walking down the tree and exploiting how the sum changes when the root shifts by one edge.

Phase 1 — Root the tree at node 0 and do a post-order DFS. For each node u compute count[u] (the size of u's subtree) and sumDist[0] for the root, which is the sum of depth[v] for all v. As we return from a child c to its parent u, the contribution is count[c] for the size and sumDist[c] + count[c] for the distance, because every node in c's subtree is one step further from u than from c.

Phase 2 — Re-root with a pre-order DFS. When the root shifts from u to a child v: every node inside v's subtree (there are count[v] of them) becomes one step closer, and every other node (there are n - count[v] of them) becomes one step farther. So sumDist[v] = sumDist[u] - count[v] + (n - count[v]). A single DFS spreads the answer to every node in O(n) total. Time is O(n) and space is O(n) for the adjacency list, counts, and answer.

Examples

Example 1
Input: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output: [8,12,6,10,10,10]
Explanation: Distances from node 0: 0+1+1+2+2+2 = 8. Distances from node 1: 1+2+2+3+3+3 = 14? The expected output uses node 2 as a hub: from node 2 the distances are 1+2+0+1+1+1 = 6, which matches the third position.
Example 2
Input: n = 1, edges = []
Output: [0]
Explanation: A single-node tree has no other nodes, so the distance sum at node 0 is 0.
Example 3
Input: n = 2, edges = [[1,0]]
Output: [1,1]
Explanation: Each of the two nodes is exactly distance 1 from the other.

Constraints

3 hints and a full solution walkthrough with step-by-step visualizations are available in the practice editor.

treedynamic-programmingdfsgraphrerooting

Browse all DSA problems