Sum of Distances in Tree: Re-rooting Technique

You are given a tree with n nodes labeled 0 to n - 1 and n - 1 edges. For each node, compute the sum of distances to every other node. Return an array where answer[i] is the sum of distances from node i to all other nodes.

This is LeetCode 834: Sum of Distances in Tree, and it is the definitive problem for learning the re-rooting technique. Instead of running BFS from every node (which would be O(n^2)), you compute the answer for one root and then shift the root to every other node in O(1) per shift, giving you O(n) total.

The problem

Given n nodes and a list of undirected edges forming a tree, return an array answer where answer[i] is the sum of distances between node i and all other nodes.

n = 6
edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]

# Output: [8, 12, 6, 10, 10, 10]
# dist from 0: 0+1+1+2+2+2 = 8
# dist from 1: 1+0+2+3+3+3 = 12
# dist from 2: 1+2+0+1+1+1 = 6
# dist from 3: 2+3+1+0+2+2 = 10
# dist from 4: 2+3+1+2+0+2 = 10
# dist from 5: 2+3+1+2+2+0 = 10

The brute force approach is to run BFS (or DFS) from every single node, summing up distances. That works, but each BFS is O(n), and you do it n times, so the total is O(n^2). For trees with up to 30,000 nodes, that is too slow.

The question is: can you compute the answer for all nodes in a single pass? Almost. You need two passes, but each is O(n).

The key insight: re-rooting

Suppose you already know dist[parent], the sum of distances from some parent node to all other nodes. Now imagine sliding the root from the parent to one of its children. What happens to the total distance?

Every node in the child's subtree gets 1 closer (they were reached through the parent, now the child is the root and they are one hop closer). Every node outside the child's subtree gets 1 farther (they used to go through the parent, now they need one extra hop through the parent).

If count[child] is the number of nodes in the child's subtree (including the child itself), then:

• count[child] nodes each get 1 closer, reducing the sum by count[child]
• n - count[child] nodes each get 1 farther, increasing the sum by n - count[child]

That gives us the re-rooting formula:

dist[child] = dist[parent] - count[child] + (n - count[child])

This is an O(1) update per edge. You walk the tree once to compute count[] and dist[root], then walk it again to propagate the answer to every node.

The solution

def sum_of_distances_in_tree(n, edges):
    graph = [[] for _ in range(n)]
    for u, v in edges:
        graph[u].append(v)
        graph[v].append(u)

    count = [1] * n
    dist = [0] * n

    def post_order(node, parent):
        for child in graph[node]:
            if child == parent:
                continue
            post_order(child, node)
            count[node] += count[child]
            dist[node] += dist[child] + count[child]

    def pre_order(node, parent):
        for child in graph[node]:
            if child == parent:
                continue
            dist[child] = dist[node] - count[child] + (n - count[child])
            pre_order(child, node)

    post_order(0, -1)
    pre_order(0, -1)
    return dist

Two DFS passes, that is the entire solution. The post-order pass computes subtree sizes and the root's total distance. The pre-order pass propagates the answer to every other node using the re-rooting formula.

Step by step walkthrough

Let's trace through the example tree with n = 6 and edges = [[0,1],[0,2],[2,3],[2,4],[2,5]].

The magic is in Step 2. Once you have dist[0] = 8, you never need to run BFS again. Each child's answer is derived from its parent's answer in O(1). Node 2 has count[2] = 4, so 4 nodes get closer and 2 nodes get farther: 8 - 4 + 2 = 6. Then from node 2, each of nodes 3, 4, 5 has count = 1, so 6 - 1 + 5 = 10.

Complexity analysis

Time is O(n). The post-order DFS visits every node once. The pre-order DFS visits every node once. Each visit does O(1) work (one addition, one subtraction).

Space is O(n) for the adjacency list, the count[] array, the dist[] array, and the recursion stack. For a balanced tree the stack is O(log n), but for a skewed tree it could be O(n).

The building blocks

This problem is built on two key techniques:

1. Re-rooting technique. When you can express the answer at a child in terms of the answer at the parent with an O(1) update, you can compute answers for all nodes in O(n) total by "re-rooting" the tree. The skeleton looks like this: compute the answer at one root with a post-order DFS, then propagate answers to all nodes with a pre-order DFS using the transition formula. This same technique appears in problems like Minimum Height Trees where you need to evaluate every node as a potential root.

2. Post-order DFS for subtree aggregation. The first pass collects subtree sizes (count[]) and accumulates the root's total distance. This is the same bottom-up pattern you see in Diameter of Binary Tree and Binary Tree Maximum Path Sum, where you gather information from children before processing the parent.

The combination of these two building blocks is what makes this problem elegant. The post-order pass gives you one correct answer. The pre-order pass turns that single answer into all answers.

Edge cases

• Single node (n = 1): return [0]. There are no other nodes, so the sum of distances is 0.
• Two nodes (n = 2): return [1, 1]. Each node is distance 1 from the other.
• Star graph (one center connected to n - 1 leaves): the center has distance n - 1. Each leaf has distance 1 + 2 * (n - 2) = 2n - 3 (distance 1 to center, distance 2 to every other leaf). The re-rooting formula handles this correctly because count[leaf] = 1 for every leaf.
• Path graph (straight line): the most skewed tree possible. The recursion stack goes O(n) deep. If n is very large, you may need to convert to an iterative approach to avoid stack overflow.
• Large n with default recursion limit: Python's default recursion limit is 1000. For trees with more than 1000 nodes, you need sys.setrecursionlimit(n + 10) or convert to iterative DFS.

From understanding to recall

You just traced through the two-pass solution and saw how a single O(1) formula turns the root's answer into every other node's answer. You understand why count[child] nodes get closer and n - count[child] nodes get farther. The formula dist[child] = dist[parent] - count[child] + (n - count[child]) makes complete sense right now.

But will you remember it in a week? Will you remember that the first pass is post-order (bottom-up to compute counts and the root's distance) and the second pass is pre-order (top-down to propagate answers)? Or will you confuse the order and compute garbage?

Understanding is not recall. Spaced repetition closes that gap. You practice reconstructing the two-pass structure from scratch at increasing intervals. After a few sessions, the re-rooting pattern becomes automatic.

The re-rooting technique and post-order subtree aggregation are two of roughly 60 reusable building blocks that cover hundreds of LeetCode problems. Drilling them individually is far more effective than grinding random problems and hoping the patterns stick.

Related posts

• Diameter of Binary Tree - Uses post-order DFS to aggregate subtree information, the same bottom-up pattern used in the first pass here.
• Binary Tree Maximum Path Sum - Another post-order tree problem where you combine left and right subtree results, illustrating how subtree aggregation works.
• Course Schedule - Graph DFS with adjacency lists and visited tracking, the same graph traversal foundation used in this problem.

Understanding re-rooting unlocks a family of tree problems where you need to evaluate every node as a potential root. Once this pattern clicks, problems that look like they require O(n^2) brute force suddenly collapse to O(n).