834. 树中距离之和

给定一个无向、连通的树。树中有 N 个标记为 0...N-1 的节点以及 N-1 条边 。

第 i 条边连接节点 edges[i][0] 和 edges[i][1] 。

返回一个表示节点 i 与其他所有节点距离之和的列表 ans。

示例 1:

输入: N = 6, edges = [[0,1],[0,2],[2,3],[2,4],[2,5]]
输出: [8,12,6,10,10,10]
解释:
如下为给定的树的示意图：
0
/ \
1 2
/|\
3 4 5

我们可以计算出 dist(0,1) + dist(0,2) + dist(0,3) + dist(0,4) + dist(0,5)
也就是 1 + 1 + 2 + 2 + 2 = 8。 因此，answer[0] = 8，以此类推。
说明: 1 <= N <= 10000

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/sum-of-distances-in-tree

 

初始化HashSet<Integer>> tree，tree[i]包含到i的所有连接节点。

初始化数组计数，count[i]对子树i中的所有节点进行计数。
初始化一个res数组，res[i]计数子树i中的距离和。

Post order dfs遍历，更新计数和res:
count[root] = sum(count[i]) + 1
res[root] = sum(res[i]) + sum(count[i])

dfs遍历，更新res:
当将根从父节点移到子节点i时，count[i] points离根节点近1,n - count[i]节点离根节点近1。
res[i] = res[根]- count[i] + N - count[i]

返回res,Done.

时间复杂度:
dfs: O (N)
dfs2: O (N)

class Solution {
    private int[] res;
    private int[] count;
    private ArrayList<HashSet<Integer>> tree;
    private int n; 
    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        tree=new ArrayList<>();
        res=new int[N];
        count=new int[N];
        n=N;
        for(int i=0;i<n;i++){
            tree.add(new HashSet<Integer>());
        }
        for(int[] eage:edges){
            tree.get(eage[0]).add(eage[1]);
            tree.get(eage[1]).add(eage[0]);
        }
        dfs(0,new HashSet<Integer>());
        dfs2(0,new HashSet<Integer>());
        return res;
    }
    private void dfs(int root,HashSet<Integer> seen){
        seen.add(root);
        for(int i:tree.get(root)){
            if(!seen.contains(i)){
                dfs(i,seen);
                count[root]+=count[i];
                res[root]+=res[i]+count[i];
            }
        }
        count[root]+=1;
    }
    private void dfs2(int root,HashSet<Integer> seen){
        seen.add(root);
        for(int i:tree.get(root)){
            if(!seen.contains(i)){
                res[i]=res[root]-count[i]+n-count[i];
                dfs2(i,seen);
            }
        }
    }
}