Sum of Distances in Tree Hard 0 attempts
LeetCode ↗

Sum of Distances in Tree

Hard TreeBinary TreeDFS LeetCode

You have an undirected tree with n nodes labeled 0 to n-1. For every node i, add up its distances to all n-1 other nodes, and return these totals as an array ans, where ans[i] holds node i's sum.

Input Format
Implement sumOfDistancesInTree(n, edges). Inputs below show JavaScript function arguments; there is no stdin parsing.
Output Format
Return the result described in the statement.

Sample Input and Output

Example 1 ยท Run Sample

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

Example 2

Input
n = 1
edges = []
Output
[0]
Constraints
  • 1 <= n <= 3 * 10^4
  • edges.length == n - 1
Test Cases (2)
Case 1
Args: [6,[[0,1],[0,2],[2,3],[2,4],[2,5]]] Expected: [8,12,6,10,10,10]
Case 2
Args: [1,[]] Expected: [0]
Topics
TreeBinary TreeDFS

Iterative Rerooting

First build a parent-before-child order. Process it backward to compute subtree sizes and distances from each subtree root. Process it forward to reroot: moving to a child makes its subtree nodes one step closer and all remaining nodes one step farther.

javascript
function sumOfDistancesInTree(n, edges) {
  const graph=Array.from({length:n},()=>[]);for(const [u,v] of edges){graph[u].push(v);graph[v].push(u);}
  const parent=Array(n).fill(-1),order=[0],count=Array(n).fill(1),answer=Array(n).fill(0);
  for(let i=0;i<order.length;i++){const u=order[i];for(const v of graph[u])if(v!==parent[u]){parent[v]=u;order.push(v);}}
  for(let i=n-1;i>0;i--){const u=order[i],p=parent[u];count[p]+=count[u];answer[p]+=answer[u]+count[u];}
  for(let i=1;i<n;i++){const u=order[i];answer[u]=answer[parent[u]]+n-2*count[u];}
  return answer;
}

Time: O(n) | Space: O(n)

Saved in this browser only. Private to you.

JavaScript