834. 树中距离之和-dfs/dp/树-困难

问题描述

给定一个无向、连通的树。树中有 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

解答

//每个节点到全局的距离，可以由父亲节点到全局的距离推出来
class Solution {
    int[][] fa;
    Map<Integer, List<Integer>> map;
    boolean[] vi;
    int[] result;
    int num;
    public void dfs(int start){
        vi[start] = true;
        List<Integer> temp = map.get(start);
        if(temp.size() == 1 && vi[temp.get(0)])return;
        for(int i:temp){
            if(vi[i])continue;
            dfs(i);
            fa[start][0] += fa[i][0]+1;
            fa[start][1] += fa[i][0]+1 + fa[i][1];
        }
    }
    public void dfs2(int start){
        vi[start] = true;
        List<Integer> temp = map.get(start);
        if(temp.size() == 1 && vi[temp.get(0)])return;
        for(int i:temp){
            if(vi[i])continue;
            //本节点到全局的距离可以由父亲节点到全局的距离推出来
            //全局距离 = 父亲节点到其他节点的距离+外部节点个数+本身作为父节点到全部子节点的距离
            result[i] = (result[start]-(fa[i][0]+fa[i][1]+1))+(num-fa[i][0]-1)+fa[i][1];
            dfs2(i);
        }
    }
    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        result = new int[N];
        vi = new boolean[N];
        if(N<=1)return result;
        map = new HashMap<>();
        num = N;
        for(int[] edge:edges){
            if(!map.containsKey(edge[0]))map.put(edge[0], new ArrayList<Integer>());
            if(!map.containsKey(edge[1]))map.put(edge[1], new ArrayList<Integer>());
            map.get(edge[0]).add(edge[1]);
            map.get(edge[1]).add(edge[0]);
        }

        fa = new int[N][2];
        //0作为起点开始dfs
        dfs(0);
        result[0] = fa[0][1];
        vi = new boolean[N];
        dfs2(0);
        return result;
    }
}