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.
這道題給了一棵樹，實際上是無向圖，讓求每個結點到其他所有結點的距離之和。這里並沒有定義樹結構，而是給了每條邊的兩端結點，那么還是先建立鄰接鏈表吧，然后當作無向圖來處理吧。由於結點的個數為N，所以直接用二維數組建立鄰接鏈表，注意無向圖是雙向的。好，現在表示樹的數據結構有了，該如何求距離之和呢？先從最簡單的例子還是看吧，假如只有一個結點的話，由於不存在其他結點，則也沒有距離之說，所以是0。若有連兩個結點，比如下面所示：
0
/
1
對於結點0來說，距離之和為1，因為只有結點1距離其為1，此時結點0只有1個子結點。若有三個結點的話，比如：
0
/ \
1 2
則所有結點到結點0的距離之和為2，而結點0也正好有兩個子結點，是不是有某種聯系呢，還是說我們想多了？再來看一個稍稍復雜些的例子吧：
0
/ \
1 2
/ \
3 4
這里的話所有結點到結點0的距離之和為6，顯然不是子結點的個數，整個樹也就5個結點。對於左子樹，這個正好是上一個討論的例子，左子樹中到結點1的距離之和為2，而左子樹總共有3個結點，加起來是5。而右子樹只有一個結點2，在右子樹中的距離之和為0，右子樹總共有1個結點，5加上1，正好是6？恭喜，這就是算每個子樹中的結點到子樹根結的距離之和的方法，即所有子結點的距離之和加上以子結點為根的子樹結點個數。說的好暈啊，用代碼來表示吧，需要兩個數組 count 和 res，其中 count[i] 表示以結點i為根結點的子樹中結點的個數，res[i] 表示其他所有結點到結點i的距離之和。根據上面的規律，可以總結出下面兩個式子：
count[root] = sum(count[i]) + 1
res[root] = sum(res[i]) + sum(count[i])
這里的 root 表示所有的子樹的根結點，i表示的是 root 的相連子結點，注意必須是相連的，這里不一定是二叉樹，所有可能會有多個子結點。另外需要注意的是這里的 res[root] 表示的是以 root 為根結點的子樹中所有的結點到 root 的距離之和，其他非子樹中結點的距離之和還沒有統計。可以發現這兩個式子中當前結點的值都是由其子結點決定的，這種由下而上的特點天然適合用后序遍歷來做，可以參見這道題 Binary Tree Postorder Traversal，還好這里不用寫迭代形式的后序遍歷，用遞歸寫就簡單的多了。同時還要注意的是用鄰接鏈表表示的無向圖遍歷時，為了避免死循環，一般是要記錄訪問過的結點的，這里由於是樹的結構，不會存在環，所以可以簡單化，直接記錄上一個結點 pre 就行了，只有當前結點i和 pre 不同才繼續處理。
好，更新完了所有的 count[root] 和 res[root] 之后，就要來更新所有的 res[i] 了，因為上面的講解提到了 res[root] 表示的是以 root 為根結點的子樹中所有的結點到 root 的距離，那么子樹之外的結點到 root 的距離也得加上，才是最終要求的 res[i]。雖然現在還沒有更新所有的 res[i]，但是有一個結點的 res 值是正確的，就是整個樹的根結點，這個真正的 res[root] 值是正確的！現在假設要計算 root 結點的一個子結點i的 res 值，即要計算所有結點到結點i的距離，此時知道以結點i為根結點的子樹的總結點個數為 count[i]，而這 count[i] 個結點之前在算 res[root] 時是到根結點 root 的距離，但是現在只要計算到結點i的距離，所以這 count[i] 個結點的距離都少了1，而其他所有的結點，共 N - count[i] 個，離結點i的距離比離 root 結點的距離都增加了1，所以 res[i] 的更新方法如下：
res[i] = res[root] - count[i] + N - count[i]
這里是從上而下的更新，可以使用最常用的先序遍歷，可以參見這道題 Binary Tree Preorder Traversal，這樣更新下來，所有的 res[i] 就都是題目中要求的值了，參見代碼如下：
class Solution {
public:
vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
vector<int> res(N), count(N);
vector<vector<int>> tree(N);
for (auto &edge : edges) {
tree[edge[0]].push_back(edge[1]);
tree[edge[1]].push_back(edge[0]);
}
helper(tree, 0, -1, count, res);
helper2(tree, 0, -1, count, res);
return res;
}
void helper(vector<vector<int>>& tree, int cur, int pre, vector<int>& count, vector<int>& res) {
for (int i : tree[cur]) {
if (i == pre) continue;
helper(tree, i, cur, count, res);
count[cur] += count[i];
res[cur] += res[i] + count[i];
}
++count[cur];
}
void helper2(vector<vector<int>>& tree, int cur, int pre, vector<int>& count, vector<int>& res) {
for (int i : tree[cur]) {
if (i == pre) continue;
res[i] = res[cur] - count[i] + count.size() - count[i];
helper2(tree, i, cur, count, res);
}
}
};
討論：整體來說，這道題算是相當有難度的一道題，同時考察了鄰接鏈表的建立，無向圖的遍歷，樹的先序和后序遍歷，以及對復雜度的拆分能力，總之是非常棒的一道題，博主非常喜歡～
Github 同步地址:
https://github.com/grandyang/leetcode/issues/834
類似題目：
Binary Tree Postorder Traversal
Binary Tree Preorder Traversal
Distribute Coins in Binary Tree
參考資料：
https://leetcode.com/problems/sum-of-distances-in-tree/
https://leetcode.com/problems/sum-of-distances-in-tree/discuss/161975/My-DFS-sulotion-two-passes
[LeetCode All in One 題目講解匯總(持續更新中...)](https://www.cnblogs.com/grandyang/p/4606334.html)