Sum of Distances in Tree

LeetCode Q 834 - Sum of Distances in Tree

An undirected, connected tree with N nodes labelled 0...N-1 and N-1 edges are given.
The ith 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]

Note: 1 <= N <= 10000

Solution:

Code:

int[] count;
int[] res;
List<HashSet<Integer>> graph;
public int[] sumOfDistancesInTree(int N, int[][] edges) {
  graph = new ArrayList<>();

  for (int i = 0; i < N; i++) graph.add(new HashSet<>());

  for (int[] e: edges) {
    graph.get(e[0]).add(e[1]);
    graph.get(e[1]).add(e[0]);
  }

  dfs(0, -1);
  dfs2(0, -1);

  return res;
}

private void dfs(int root, int parent) {
  for (int nei: graph.get(root)) {
    if (nei == parent) continue;
    dfs(nei, root);
    count[root] += count[nei];
    res[root] += res[nei] + count[nei];
  }
  count[root]++;
}
    
private void dfs2(int root, int parent) {
  for (int nei: graph.get(root)) {
    if (nei == parent) continue;
    res[nei] = res[root] + (count.length - count[nei]) - count[nei];
    dfs2(nei, root);
  }
}