Sum of Distances in Tree — LeetCode 834 Rerooting Masterclass

Problem Statement#

You are given an undirected, connected tree with n nodes labeled 0..n-1 and n-1 edges. Return an array answer where answer[i] is the sum of distances between node i and every other node in the tree.

Constraints:

• 1 <= n <= 3 * 10^4
• edges.length == n - 1
• The given input represents a valid tree

Input:  n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output: [8,12,6,10,10,10]

Input:  n = 1, edges = []
Output: [0]

Why This Problem Matters#

LeetCode 834 — Sum of Distances in Tree — is a hard graph problem that frequently appears in interviews at Google, Meta, Amazon, and Microsoft. A brute-force solution runs BFS or DFS from every node giving O(n^2) which times out for n = 3 * 10^4. Interviewers love this one because it forces candidates to think about reusing subtree work via the rerooting (or "switch root") DP technique on trees.

This is a defining problem for the FAANG tree-DP toolkit. Recognizing that the answer for a child can be derived from the parent's answer in O(1) is the kind of mental leap interviewers reward. The same rerooting pattern also solves LC 310 Minimum Height Trees, LC 968 Binary Tree Cameras, and competitive programming staples.

The pattern transfers directly to social-network distance computations, infection spread modeling, and routing in tree-shaped topologies — making this problem highly relevant to backend, distributed-systems, and ML infra teams at FAANG.

The Core Insight#

Root the tree at node 0. For each node u, let count[u] be the size of its subtree and dist[u] be the sum of distances from u to every node inside its subtree. Both can be computed in one post-order DFS.

The magic step is rerooting. If we already know answer[parent], moving the root from parent to child:

• All count[child] nodes inside child's subtree get one step closer, decreasing the sum by count[child].
• All n - count[child] nodes outside child's subtree get one step farther, increasing the sum by n - count[child].

So answer[child] = answer[parent] - count[child] + (n - count[child]). A second pre-order DFS propagates this from the root downward in O(n).

Visual Dry Run#

Tree: 0 - 1, 0 - 2, 2 - 3, 2 - 4, 2 - 5. Root at 0.

Pre-order rerooting from answer[0] = 8:

• answer[1] = 8 - 1 + (6 - 1) = 12
• answer[2] = 8 - 4 + (6 - 4) = 6
• answer[3] = 6 - 1 + (6 - 1) = 10

Solution (Optimal)#

from typing import List
from collections import defaultdict
 
class Solution:
    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]:
        graph = defaultdict(list)
        for a, b in edges:
            graph[a].append(b)
            graph[b].append(a)
 
        count = [1] * n
        answer = [0] * n
 
        def post_order(node: int, parent: int) -> None:
            for nei in graph[node]:
                if nei == parent:
                    continue
                post_order(nei, node)
                count[node] += count[nei]
                answer[node] += answer[nei] + count[nei]
 
        def pre_order(node: int, parent: int) -> None:
            for nei in graph[node]:
                if nei == parent:
                    continue
                answer[nei] = answer[node] - count[nei] + (n - count[nei])
                pre_order(nei, node)
 
        post_order(0, -1)
        pre_order(0, -1)
        return answer

var sumOfDistancesInTree = function(n, edges) {
    const graph = Array.from({ length: n }, () => []);
    for (const [a, b] of edges) {
        graph[a].push(b);
        graph[b].push(a);
    }
    const count = new Array(n).fill(1);
    const answer = new Array(n).fill(0);
 
    const postOrder = (node, parent) => {
        for (const nei of graph[node]) {
            if (nei === parent) continue;
            postOrder(nei, node);
            count[node] += count[nei];
            answer[node] += answer[nei] + count[nei];
        }
    };
 
    const preOrder = (node, parent) => {
        for (const nei of graph[node]) {
            if (nei === parent) continue;
            answer[nei] = answer[node] - count[nei] + (n - count[nei]);
            preOrder(nei, node);
        }
    };
 
    postOrder(0, -1);
    preOrder(0, -1);
    return answer;
};

Time: O(n) — two DFS traversals, each visiting every edge once. Space: O(n) — adjacency list, count and answer arrays, recursion stack.

Common Mistakes#

• Computing dist[root] using only depths — works but most candidates conflate "sum of distances" with "sum of depths" and forget to derive other roots.
• Forgetting to skip the parent in the adjacency list, leading to infinite recursion.
• Using BFS with O(n) queue rebuild per node (O(n^2)) and TLE-ing.
• Iterating once and trying to compute answer[child] before answer[parent] is finalized — order matters.
• Recursion depth: with n = 30000, deep skewed trees can blow Python's default recursion limit. Use sys.setrecursionlimit or convert to iterative.

Interview Tips#

• Draw a 5-node tree, root it, and verbally walk through both DFS passes.
• State the rerooting formula explicitly: answer[child] = answer[parent] - count[child] + (n - count[child]).
• Mention that this is the same family as LC 310 Minimum Height Trees and LC 1483 Kth Ancestor.
• If the interviewer asks for a brute force first, give O(n^2) BFS-from-every-node, then optimize.

Follow-up Questions#

• "What if the tree is weighted?" — Replace + count[nei] with + count[nei] * w(node, nei) and adjust the rerooting delta accordingly.
• "Sum of squared distances?" — Need a second DP storing sum of squared depths.
• "Online queries: distance between u and v?" — Use binary lifting for LCA in O(log n) per query.
• "Diameter using rerooting?" — Track two longest down-paths per node, classic O(n) tree DP.
• "Handle dynamic edge deletions?" — Falls outside basic rerooting; use Euler tour + segment tree.

Key Takeaways#

• LeetCode 834 Sum of Distances in Tree is a hard tree-DP problem solvable in O(n) using rerooting.
• Two DFS passes: post-order computes subtree size and within-subtree sum; pre-order rerooting propagates the global answer.
• The rerooting formula is answer[child] = answer[parent] - count[child] + (n - count[child]).
• Brute force runs O(n^2) and times out for n = 30000.
• Companies like Google, Meta, and Amazon use this pattern for graph-distance interviews.
• The technique generalizes to weighted trees, diameter, and any DP that depends on the root choice.
• Watch recursion depth on skewed trees — Python's default limit is 1000.