Sum of Distances in Tree Solution

Problem

Given a tree of n nodes labeled 0 through n−1, find the sum of distances from each node to every other node. Return an array where answer[i] is the total number of edges you'd need to traverse to visit all other nodes starting from node i.

For a tree where node 0 connects to nodes 1 and 2, and node 1 connects to nodes 3 and 4:

[0,1,2,3,4]

Input: n = 5, edges = [[0,1],[0,2],[1,3],[1,4]]
Output: [6, 5, 9, 8, 8]
Explanation: Node 1 is centrally positioned — one hop from its parent and one hop from each of its two children — giving it the smallest total distance; node 2 is a leaf that must cross the entire tree to reach nodes 3 and 4.

Intuition

The brute-force approach — run BFS from every node — is O(n²), which breaks at n = 30,000. The key insight is that once you know the total distance for one node, you can derive the answer for any neighbor in O(1). When you "move the root" from a parent to a child, the count[child] nodes in the child's subtree each get one step closer while the remaining n − count[child] nodes each get one step farther. A single top-down pass propagates every answer from the root outward.

Approach 1 — Brute Force BFS

For each node, run a full BFS to measure the sum of distances to all other nodes.

Build an undirected adjacency list from the edges.
For each starting node, run BFS while tracking the distance from that start.
Accumulate the running total and store in the answer array.

1from collections import deque, defaultdict
2
3def sumOfDistancesInTree(n: int, edges: list[list[int]]) -> list[int]:
4    graph = defaultdict(list)
5    for u, v in edges:
6        graph[u].append(v)
7        graph[v].append(u)
8
9    def bfs_from(start: int) -> int:
10        visited = {start}
11        queue = deque([(start, 0)])
12        total_distance = 0
13        while queue:
14            node, dist = queue.popleft()
15            total_distance += dist
16            for neighbor in graph[node]:
17                if neighbor not in visited:
18                    visited.add(neighbor)
19                    queue.append((neighbor, dist + 1))
20        return total_distance
21
22    return [bfs_from(i) for i in range(n)]

Time: O(n²) — n separate BFS traversals each visiting all n nodes.

Space: O(n) — visited array and queue reused each BFS run.

Approach 2 — Two-Pass BFS (Re-rooting)

Root the tree at node 0, compute the root's answer bottom-up, then derive every other answer top-down in one additional pass.

BFS from node 0 to record a traversal order and each node's parent.
Bottom-up (reverse BFS order): for each node, roll up count[node] (subtree size including self) and dp[node] (total distance to every node within that subtree).
After this pass, dp[0] is the correct answer for the root.
Top-down (BFS order): for each child derive its answer from its parent using the re-rooting formula: dp[child] = dp[parent] - count[child] + (n - count[child]).
Return dp as the final answer.

1from collections import deque, defaultdict
2
3def sumOfDistancesInTree(n: int, edges: list[list[int]]) -> list[int]:
4    graph = defaultdict(list)
5    for u, v in edges:
6        graph[u].append(v)
7        graph[v].append(u)
8
9    count = [1] * n   # subtree size rooted at each node (starts at 1 for self)
10    dp = [0] * n      # sum of distances to all nodes within the rooted subtree
11    parent = [-1] * n
12    visited = [False] * n
13    order = []
14
15    # BFS to fix a parent-child relationship and processing order
16    queue = deque([0])
17    visited[0] = True
18    while queue:
19        node = queue.popleft()
20        order.append(node)
21        for neighbor in graph[node]:
22            if not visited[neighbor]:
23                visited[neighbor] = True
24                parent[neighbor] = node
25                queue.append(neighbor)
26
27    # Bottom-up: leaves first, accumulate subtree sizes and intra-subtree distances
28    for node in reversed(order):
29        if parent[node] != -1:
30            count[parent[node]] += count[node]
31            # every node in child's subtree is one extra edge away from parent
32            dp[parent[node]] += dp[node] + count[node]
33
34    # Top-down: re-root — count[child] nodes get closer, the rest get farther
35    for node in order[1:]:
36        par = parent[node]
37        dp[node] = dp[par] - count[node] + (n - count[node])
38
39    return dp

Time: O(n) — one BFS pass to build order, then two linear sweeps over the n nodes.

Space: O(n) — adjacency list, count, dp, parent, and order arrays each of size n.

Complexity Summary

Approach	Time	Space	When to use
Brute Force BFS	O(n²)	O(n)	Small trees (n ≤ 1,000) or as a reference implementation
Two-Pass BFS (Re-rooting)	O(n)	O(n)	Any real case; handles the full n = 30,000 constraint

Common Mistakes

Missing + count[child] in the bottom-up step: Writing dp[parent] += dp[child] alone ignores the direct edge between parent and child — every node in the child's subtree is one additional hop away from the parent.
Inverted re-rooting sign: Writing dp[child] = dp[parent] + count[child] - (n - count[child]) swaps which group moves closer and which moves farther — the subtree nodes get closer (subtract), the outside nodes get farther (add), not the reverse.
Not guarding against cycling: In the recursive DFS version of this approach, failing to skip the parent neighbor causes infinite recursion. The iterative BFS version avoids this naturally since visited nodes are marked.
Initializing count to 0: Each node is in its own subtree, so count must start at [1] * n. Starting at 0 makes the root's subtree count off by 1 and cascades incorrect distances through the whole tree.
Python recursion limit: A linear chain of n = 30,000 nodes reaches 30,000 frames deep, crashing the default recursive DFS. The iterative BFS-based approach shown here has no such limit.

Sum of Distances in Tree — Solution