Sum of Distances in Tree

Hard

#0788Time: O(N^2). For each of the N nodes, we perform a BFS traversal. A single BFS on a tree takes O(N + E) time, where E is the number of edges (N-1). So, one BFS is O(N). Repeating this for all N nodes results in a total time complexity of O(N * N) = O(N^2).Space: O(N). The adjacency list requires O(N + E) = O(N) space. The queue and the `distances` array used in each BFS also require O(N) space.4 companies

Patterns

Dynamic Programming

Algorithms

Depth-First Search

Data structures

Tree Graph

Companies

ByteDance Media.net PhonePe MathWorks

Prompt

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 the i^th node in the tree and all other nodes.

Example 1:

Input: n = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
Output: [8,12,6,10,10,10]
Explanation: The tree is shown above.
We can see that dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
equals 1 + 1 + 2 + 2 + 2 = 8.
Hence, answer[0] = 8, and so on.

Example 2:

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

Example 3:

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

Constraints:

1 <= n <= 3 * 10⁴
edges.length == n - 1
edges[i].length == 2
0 <= a_i, b_i < n
a_i != b_i
The given input represents a valid tree.

Approaches

2 approaches with complexity analysis and trade-offs.

This approach is the most straightforward. For every single node in the tree, we can calculate the sum of its distances to all other nodes. To find the distance from a starting node u to all other nodes v, we can use a graph traversal algorithm like Breadth-First Search (BFS) or Depth-First Search (DFS). BFS is a natural choice here as it explores the tree layer by layer, directly giving the shortest distance from the source.

Algorithm

1. Create an adjacency list adj to represent the tree.
1. Initialize an answer array of size n.
1. For each node i from 0 to n-1:
- a. Initialize a queue for BFS and add i to it.
- b. Initialize a distances array of size n (e.g., filled with -1) and set distances[i] = 0.
- c. Initialize a variable currentSum = 0.
- d. While the queue is not empty:
  - i. Dequeue a node u.
  - ii. Add distances[u] to currentSum.
  - iii. For each neighbor v of u:
    - If v has not been visited (distances[v] == -1):
      - Set distances[v] = distances[u] + 1.
      - Enqueue v.
- e. Store the result: answer[i] = currentSum.
1. Return the answer array.

Walkthrough

We iterate through each node i from 0 to n-1. For each i, we perform a BFS starting from it to compute the distances to all other nodes. A queue is used for the BFS, storing pairs of (node, distance). We also need a visited array for each BFS run to avoid re-visiting nodes. During the BFS starting from i, we maintain a running sum. When we visit a node j at a distance d, we add d to this sum. After the BFS completes, the total sum is the answer for node i, so we store it in answer[i]. This process is repeated for all n nodes. The first step is to build an adjacency list representation of the tree from the input edges array to facilitate the traversal.

1import java.util.*;2 3class Solution {4    public int[] sumOfDistancesInTree(int n, int[][] edges) {5        if (n == 1) {6            return new int[]{0};7        }8 9        List<List<Integer>> adj = new ArrayList<>();10        for (int i = 0; i < n; i++) {11            adj.add(new ArrayList<>());12        }13        for (int[] edge : edges) {14            adj.get(edge[0]).add(edge[1]);15            adj.get(edge[1]).add(edge[0]);16        }17 18        int[] answer = new int[n];19        for (int i = 0; i < n; i++) {20            answer[i] = bfs(i, n, adj);21        }22        return answer;23    }24 25    private int bfs(int startNode, int n, List<List<Integer>> adj) {26        int[] dist = new int[n];27        Arrays.fill(dist, -1);28        Queue<Integer> queue = new LinkedList<>();29 30        queue.offer(startNode);31        dist[startNode] = 0;32        int totalDistance = 0;33 34        while (!queue.isEmpty()) {35            int u = queue.poll();36            totalDistance += dist[u];37 38            for (int v : adj.get(u)) {39                if (dist[v] == -1) {40                    dist[v] = dist[u] + 1;41                    queue.offer(v);42                }43            }44        }45        return totalDistance;46    }47}

Complexity

Time

O(N^2). For each of the N nodes, we perform a BFS traversal. A single BFS on a tree takes O(N + E) time, where E is the number of edges (N-1). So, one BFS is O(N). Repeating this for all N nodes results in a total time complexity of O(N * N) = O(N^2).

Space

O(N). The adjacency list requires O(N + E) = O(N) space. The queue and the `distances` array used in each BFS also require O(N) space.

Trade-offs

Pros

Conceptually simple and easy to implement.
It's a direct translation of the problem statement.

Cons

Inefficient. The O(N^2) time complexity will result in a "Time Limit Exceeded" (TLE) for larger inputs (like N = 3 * 10^4).
It performs a lot of redundant calculations. The distance between two nodes u and v is calculated twice, once when starting from u and once when starting from v.

A more efficient solution can be achieved by recognizing the relationship between the answers for adjacent nodes. Instead of re-calculating everything from scratch for each node, we can derive the answer for a node based on the answer of its parent. This dynamic programming approach on the tree requires two Depth-First Search (DFS) traversals.

Algorithm

1. Build an adjacency list adj for the tree.
1. Initialize two arrays, count (for subtree sizes) and ans (for the final answer), both of size n.
1. First Pass (post-order DFS, e.g., dfs1(u, parent)):
- a. Start DFS from an arbitrary root (e.g., node 0).
- b. For a node u, recursively call dfs1 for all its children v.
- c. After the children's calls return, update count[u] and ans[u]:
  - count[u] = 1 + Σ count[v] for all children v.
  - ans[u] = Σ (ans[v] + count[v]) for all children v.
1. Second Pass (pre-order DFS, e.g., dfs2(u, parent)):
- a. Start DFS from the same root (node 0).
- b. For a node u, before visiting its children, its ans[u] is already correct.
- c. For each child v of u, calculate its final answer using the parent's answer:
  - ans[v] = ans[u] - count[v] + (n - count[v]).
- d. Recursively call dfs2 for the child v.
1. Return the ans array.

Walkthrough

The core idea is to first calculate the answer for an arbitrary root (say, node 0) and then use this result to efficiently calculate the answers for all other nodes. This is done in two passes.

Pass 1: Post-order Traversal (Bottom-up) We perform a DFS starting from root 0. This traversal will compute two things for each node u:

count[u]: The number of nodes in the subtree rooted at u (including u).
ans[u]: The sum of distances from u to all other nodes within its own subtree. The DFS explores children first. After returning from all children v of u, we can compute count[u] and ans[u] for the parent u.
count[u] = 1 + sum(count[v]) for all children v.
ans[u] = sum(ans[v] + count[v]) for all children v. The term ans[v] is the sum of distances from v to its descendants. To get the sum of distances from u to v's descendants, we note that each path is one edge longer (the edge u-v), so we add count[v] (the number of such descendants). After this pass, ans[0] holds the correct final answer for the root node, as its subtree is the entire tree.

Pass 2: Pre-order Traversal (Top-down) We perform a second DFS, also starting from root 0. This traversal will use the parent's correct answer to calculate the child's correct answer. When we move from a parent u to a child v, the sum of distances changes.

For the count[v] nodes in v's subtree, their distance to the new "root" v becomes 1 shorter. This contributes a change of -count[v].
For the n - count[v] nodes outside v's subtree, their distance to the new "root" v becomes 1 longer. This contributes a change of +(n - count[v]). So, the update formula is: ans[v] = ans[u] - count[v] + (n - count[v]). We apply this formula as we traverse down the tree, and after this pass, the ans array is fully populated with the correct values.

1import java.util.*;2 3class Solution {4    private List<List<Integer>> adj;5    private int[] count;6    private int[] ans;7    private int n;8 9    public int[] sumOfDistancesInTree(int n, int[][] edges) {10        this.n = n;11        this.adj = new ArrayList<>();12        for (int i = 0; i < n; i++) {13            adj.add(new ArrayList<>());14        }15        for (int[] edge : edges) {16            adj.get(edge[0]).add(edge[1]);17            adj.get(edge[1]).add(edge[0]);18        }19 20        this.count = new int[n];21        this.ans = new int[n];22 23        // Pass 1: Post-order traversal to calculate subtree sizes and partial sums24        dfs1(0, -1);25        26        // Pass 2: Pre-order traversal to update sums for all nodes27        dfs2(0, -1);28 29        return ans;30    }31 32    // Post-order traversal: calculate count[u] and ans[u] for the subtree rooted at u33    private void dfs1(int u, int parent) {34        count[u] = 1;35        for (int v : adj.get(u)) {36            if (v == parent) continue;37            dfs1(v, u);38            count[u] += count[v];39            ans[u] += ans[v] + count[v];40        }41    }42 43    // Pre-order traversal: calculate the final answer for each node44    private void dfs2(int u, int parent) {45        for (int v : adj.get(u)) {46            if (v == parent) continue;47            // ans[v] is currently the sum of distances in its subtree.48            // We need to update it based on the parent's final answer.49            // ans[u] is the final answer for node u.50            // Move root from u to v:51            // - nodes in v's subtree get 1 closer: -count[v]52            // - nodes outside v's subtree get 1 farther: +(n - count[v])53            ans[v] = ans[u] - count[v] + (n - count[v]);54            dfs2(v, u);55        }56    }57}

Complexity

Time

O(N). We perform two separate DFS traversals on the tree. Each traversal visits every node and edge exactly once. Building the adjacency list also takes O(N) time. Thus, the total time complexity is linear with respect to the number of nodes.

Space

O(N). The adjacency list requires O(N) space. The `count` and `ans` arrays also take O(N) space. The recursion stack for DFS can be up to O(N) in the worst case (for a skewed tree), leading to an overall O(N) space complexity.

Trade-offs

Pros

Highly efficient with linear time complexity, making it suitable for large inputs.
It's the optimal way to solve this problem.

Cons

The logic is more complex than the brute-force approach.
It requires a good understanding of tree traversals and dynamic programming on trees.

Solutions

Solution

1class Solution {2private3  int n;4private5  int[] ans;6private7  int[] size;8private9  List<Integer>[] g;10public11  int[] sumOfDistancesInTree(int n, int[][] edges) {12    this.n = n;13    g = new List[n];14    ans = new int[n];15    size = new int[n];16    Arrays.setAll(g, k->new ArrayList<>());17    for (var e : edges) {18      int a = e[0], b = e[1];19      g[a].add(b);20      g[b].add(a);21    }22    dfs1(0, -1, 0);23    dfs2(0, -1, ans[0]);24    return ans;25  }26private27  void dfs1(int i, int fa, int d) {28    ans[0] += d;29    size[i] = 1;30    for (int j : g[i]) {31      if (j != fa) {32        dfs1(j, i, d + 1);33        size[i] += size[j];34      }35    }36  }37private38  void dfs2(int i, int fa, int t) {39    ans[i] = t;40    for (int j : g[i]) {41      if (j != fa) {42        dfs2(j, i, t - size[j] + n - size[j]);43      }44    }45  }46}

1class Solution {2public:3  vector<int> sumOfDistancesInTree(int n, vector<vector<int>> &edges) {4    vector<vector<int>> g(n);5    for (auto &e : edges) {6      int a = e[0], b = e[1];7      g[a].push_back(b);8      g[b].push_back(a);9    }10    vector<int> ans(n);11    vector<int> size(n);12    function<void(int, int, int)> dfs1 = [&](int i, int fa, int d) {13      ans[0] += d;14      size[i] = 1;15      for (int &j : g[i]) {16        if (j != fa) {17          dfs1(j, i, d + 1);18          size[i] += size[j];19        }20      }21    };22    function<void(int, int, int)> dfs2 = [&](int i, int fa, int t) {23      ans[i] = t;24      for (int &j : g[i]) {25        if (j != fa) {26          dfs2(j, i, t - size[j] + n - size[j]);27        }28      }29    };30    dfs1(0, -1, 0);31    dfs2(0, -1, ans[0]);32    return ans;33  }34};

1class Solution:2    def sumOfDistancesInTree(self, n: int, edges: List[List[int]]) -> List[int]: def dfs1(i: int, fa: int, d: int): ans[0] += d size[i] = 1 for j in g[i]: if j != fa: dfs1(j, i, d + 1) size[i] += size[j] def dfs2(i: int, fa: int, t: int): ans[i] = t for j in g[i]: if j != fa: dfs2(j, i, t - size[j] + n - size[j]) g = defaultdict(list) for a, b in edges: g[a]. append(b) g[b]. append(a) ans = [0] * n size = [0] * n dfs1(0, - 1, 0) dfs2(0, - 1, ans[0]) return ans

Video walkthrough

Newsletter

One sharp idea, every week

System design and interview prep — short enough to finish.

No spam. Unsubscribe anytime.

Practice

Same difficulty — related problems to reinforce the pattern.

← Previous

Previous problem

Next problem