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]
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 1:

BFS

class Solution {
    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        List<Integer>[] list = new List[N];
        for (int i = 0; i < N; i ++) {
            list[i] = new ArrayList();
        }
        for (int[] edge : edges) {
            list[edge[0]].add(edge[1]);
            list[edge[1]].add(edge[0]);
        }
        int[] result = new int[N];
        for (int i = 0; i < N; i ++) {
            Deque<Integer> queue = new ArrayDeque();
            boolean[] visited = new boolean[N];
            visited[i] = true;
            queue.offer(i);
            int step = 0;
            while (!queue.isEmpty()) {
                int size = queue.size();
                step ++;
                for (int j = 0; j < size; j ++) {
                    int curr = queue.poll();
                    for (int next : list[curr]) {
                        if (!visited[next]) {
                            result[i] += step;
                            visited[next] = true;
                            queue.offer(next);
                        }
                    }
                }
            }
        }
        return result;
    }
}

solution 2:

BFS

https://leetcode.com/problems/sum-of-distances-in-tree/discuss/130583/C%2B%2BJavaPython-Pre-order-and-Post-order-DFS-O(N)

class Solution {
    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        Set<Integer>[] tree = new Set[N];
        for (int i = 0; i < N; i ++) {
            tree[i] = new HashSet();
        }
        for (int[] edge : edges) {
            tree[edge[0]].add(edge[1]);
            tree[edge[1]].add(edge[0]);
        }
        int[] result = new int[N];
        int[] count = new int[N];
        
        post(tree, count, result, 0, -1);
        pre(tree, count, result, 0, -1);
        
        return result;
    }
    
    private void post(Set<Integer>[] tree, int[] count, int[] result, int root, int prev) {
        for (int child : tree[root]) {
            if (child == prev) {
                continue;
            }
            post(tree, count, result, child, root);
            count[root] += count[child];
            result[root] += result[child] + count[child];
        }
        count[root] ++;
    }
    
    private void pre(Set<Integer>[] tree, int[] count, int[] result, int root, int prev) {
        for (int child : tree[root]) {
            if (child == prev) {
                continue;
            }
            result[child] = result[root] - count[child] + (count.length - count[child]);
            pre(tree, count, result, child, root);
        }
    }
}