Sum of Distances in Tree
HardPrompt
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] = [ai, bi] indicates that there is an edge between nodes ai and bi in the tree.
Return an array answer of length n where answer[i] is the sum of the distances between the ith 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 * 104edges.length == n - 1edges[i].length == 20 <= ai, bi < nai != bi- 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
-
- Create an adjacency list
adjto represent the tree.
- Create an adjacency list
-
- Initialize an
answerarray of sizen.
- Initialize an
-
- For each node
ifrom0ton-1:
- a. Initialize a queue for BFS and add
ito it. - b. Initialize a
distancesarray of sizen(e.g., filled with -1) and setdistances[i] = 0. - c. Initialize a variable
currentSum = 0. - d. While the queue is not empty:
- i. Dequeue a node
u. - ii. Add
distances[u]tocurrentSum. - iii. For each neighbor
vofu:- If
vhas not been visited (distances[v] == -1):- Set
distances[v] = distances[u] + 1. - Enqueue
v.
- Set
- If
- i. Dequeue a node
- e. Store the result:
answer[i] = currentSum.
- For each node
-
- Return the
answerarray.
- Return the
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.
import java.util.*; class Solution { public int[] sumOfDistancesInTree(int n, int[][] edges) { if (n == 1) { return new int[]{0}; } List<List<Integer>> adj = new ArrayList<>(); for (int i = 0; i < n; i++) { adj.add(new ArrayList<>()); } for (int[] edge : edges) { adj.get(edge[0]).add(edge[1]); adj.get(edge[1]).add(edge[0]); } int[] answer = new int[n]; for (int i = 0; i < n; i++) { answer[i] = bfs(i, n, adj); } return answer; } private int bfs(int startNode, int n, List<List<Integer>> adj) { int[] dist = new int[n]; Arrays.fill(dist, -1); Queue<Integer> queue = new LinkedList<>(); queue.offer(startNode); dist[startNode] = 0; int totalDistance = 0; while (!queue.isEmpty()) { int u = queue.poll(); totalDistance += dist[u]; for (int v : adj.get(u)) { if (dist[v] == -1) { dist[v] = dist[u] + 1; queue.offer(v); } } } return totalDistance; }}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
uandvis calculated twice, once when starting fromuand once when starting fromv.
Solutions
Solution
class Solution {private int n;private int[] ans;private int[] size;private List<Integer>[] g;public int[] sumOfDistancesInTree(int n, int[][] edges) { this.n = n; g = new List[n]; ans = new int[n]; size = new int[n]; Arrays.setAll(g, k->new ArrayList<>()); for (var e : edges) { int a = e[0], b = e[1]; g[a].add(b); g[b].add(a); } dfs1(0, -1, 0); dfs2(0, -1, ans[0]); return ans; }private void dfs1(int i, int fa, int d) { ans[0] += d; size[i] = 1; for (int j : g[i]) { if (j != fa) { dfs1(j, i, d + 1); size[i] += size[j]; } } }private void dfs2(int i, int fa, int t) { ans[i] = t; for (int j : g[i]) { if (j != fa) { dfs2(j, i, t - size[j] + n - size[j]); } } }}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.