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;
        }
    }
}