Sum of Distances in Tree

Question

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.

LeetCodeleetcode.com/problems/sum-of-distances-in-tree

Solution

/*
@param: int[][] edges
@return: int[]
Algorithm: 第一遍dfs先把count[] 和 child[] 计算出来
count[i]: i节点subtree的sum， child[i]: i节点的subtree的节点数
第二遍dfs 计算每一个点的sum
当从i走到j时，有child[j]个节点的sum-1，N-child[j]个节点sum+1；
count[j] = count[i] - child[j] + N - child[j];
*/
int[] count, child;
int l;
public int[] sumOfDistancesInTree(int N, int[][] edges) {        
    Map<Integer,List<Integer>> neighbors = new HashMap<>();
    for (int[] n : edges) {
        int a = n[0], b = n[1];
        neighbors.computeIfAbsent(a, k -> new ArrayList<>());
        neighbors.computeIfAbsent(b, k -> new ArrayList<>());
        neighbors.get(a).add(b);
        neighbors.get(b).add(a);
    }
    
    l = N;
    child = new int[N];
    count = new int[N];
    Arrays.fill(child, 1);
    
    if (neighbors.containsKey(0)){
        dfs1(0, -1, neighbors);
        dfs2(0, -1, neighbors);            
    }

    return count;
}

private void dfs1(int cur, int pre, Map<Integer,List<Integer>> neighbors) {
    for (int nxt : neighbors.get(cur)) {
        if (nxt == pre) continue;
        dfs1(nxt, cur, neighbors);
        child[cur] += child[nxt];
        count[cur] += count[nxt] + child[nxt];
    }        
}

private void dfs2(int cur, int pre, Map<Integer,List<Integer>> neighbors) {
    for (int nxt : neighbors.get(cur)) {
        if (nxt == pre) continue;
        count[nxt] = count[cur] - child[nxt] + l - child[nxt];
        dfs2(nxt, cur, neighbors);
    }        
}