Formatted question description: https://leetcode.ca/all/834.html

# 834. Sum of Distances in Tree

## Level

Hard

## Description

An undirected, connected tree with `N`

nodes labelled `0...N-1`

and `N-1`

`edges`

are given.

The `i`

th edge connects nodes `edges[i][0]`

and `edges[i][1]`

together.

Return a list `ans`

, where `ans[i]`

is the sum of the distances between node `i`

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:
Here is a diagram of the given tree:
0
/ \
1 2
/|\
3 4 5
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.
```

Note: `1 <= N <= 10000`

## Solution

Create two arrays `counts`

and `sums`

of length `N`

, where `counts[i]`

represents the number of nodes in the subtree that has root `i`

and `sums`

is the result array. Initialize all elements in `counts`

to 1.

Do depth first search starting from node 0 for two times.

For the first time, calculate the array `counts`

and each root’s value of `sums`

. For each node `i`

that is not a leaf, `counts[i]`

is the sum of all elements `counts[j]`

plus 1 and `sums[i]`

is the sum of all elements `sums[j]`

and `counts[j]`

, where `j`

is a child of `i`

.

For the second time, calculate each child’s value of `sums`

. For each node `i`

that is not 0, if its parent is `j`

, then `sums[i] = sums[j] + counts[i] + N - counts[i]`

.

```
class Solution {
public int[] sumOfDistancesInTree(int N, int[][] edges) {
Map<Integer, Set<Integer>> edgesMap = new HashMap<Integer, Set<Integer>>();
for (int[] edge : edges) {
int node1 = edge[0], node2 = edge[1];
Set<Integer> set1 = edgesMap.getOrDefault(node1, new HashSet<Integer>());
Set<Integer> set2 = edgesMap.getOrDefault(node2, new HashSet<Integer>());
set1.add(node2);
set2.add(node1);
edgesMap.put(node1, set1);
edgesMap.put(node2, set2);
}
int[] counts = new int[N];
Arrays.fill(counts, 1);
int[] sums = new int[N];
depthFirstSearch1(0, -1, edgesMap, counts, sums);
depthFirstSearch2(0, -1, edgesMap, counts, sums, N);
return sums;
}
public void depthFirstSearch1(int node, int parent, Map<Integer, Set<Integer>> edgesMap, int[] counts, int[] sums) {
Set<Integer> children = edgesMap.getOrDefault(node, new HashSet<Integer>());
for (int child : children) {
if (child != parent) {
depthFirstSearch1(child, node, edgesMap, counts, sums);
counts[node] += counts[child];
sums[node] += sums[child] + counts[child];
}
}
}
public void depthFirstSearch2(int node, int parent, Map<Integer, Set<Integer>> edgesMap, int[] counts, int[] sums, int length) {
Set<Integer> children = edgesMap.getOrDefault(node, new HashSet<Integer>());
for (int child : children) {
if (child != parent) {
sums[child] = sums[node] - counts[child] + length - counts[child];
depthFirstSearch2(child, node, edgesMap, counts, sums, length);
}
}
}
}
```