834. 树中距离之和-困难-树、图、动态规划、深度优先搜索

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

题解

这道题考察了很多的知识点，不过首先还是要想到动态规划，而此题的难点也在于动态规划的公式（即节点变化时，该如何计算）

复习时建议先自己好好思考一下，想不出来可以看一下提示，还是想不出来就看一看官网的题解——834. 树中距离之和

class Solution {

    int ans[];  // 结果储存
    int sz[];  // 存储树的节点数
    int dp[];  // 动态规划存储
    List<List<Integer>> graph;  // 把树变为图

    public int[] sumOfDistancesInTree(int N, int[][] edges) {
        graph = new ArrayList<List<Integer>>();
        ans = new int[N];
        sz = new int[N];
        dp = new int[N];

        for(int i=0; i<N; i++){
            graph.add(new ArrayList<Integer>());
        }

        for(int[] edge:edges){
            int u = edge[0];
            int v = edge[1];

            graph.get(u).add(v);
            graph.get(v).add(u);
        }
        
        DFS(0, -1);
        DFS2(0, -1);

        return ans;
    }

    public void DFS(int child, int parent){
        // 初始化
        sz[child] = 1;
        dp[child] = 0;
        // 遍历连接当前节点的其他节点
        for( int node:graph.get(child) ){
            if(node == parent) continue;
            DFS(node, child);
            dp[child] += dp[node] + sz[node];
            sz[child] += sz[node];
        }
    }

    // 由于若每一个节点都进行一次DFS,则将超时，故需要通过动态规划方法，降低复杂度
    public void DFS2(int child, int parent){
        ans[child] = dp[child];

        for(int node:graph.get(child)){
            if(node == parent) continue;

            // 将值存储起来
            int dp_child = dp[child];  
            int sz_child = sz[child];
            int dp_node = dp[node];
            int sz_node = sz[node];

            // 根据公式，求出dp[node]
            dp[child] -= dp[node] + sz[node];
            sz[child] -= sz[node];
            dp[node] += dp[child] + sz[child];
            sz[node] += sz[child];

            DFS2(node, child);  // 将求出的结果赋值给ans

            // 每次结尾将dp[child]等值变回原来的值,以方便其他节点进行调用求解
            dp[child] = dp_child;
            sz[child] = sz_child;
            dp[node] = dp_node;
            sz[node] = sz_node;
        }
    }
}