[LeetCode] 834. Sum of Distances in Tree 樹中距離之和

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

https://leetcode.com/problems/sum-of-distances-in-tree/discuss/130583/C%2B%2BJavaPython-Pre-order-and-Post-order-DFS-O(N)

LeetCode All in One 題目講解彙總(持續更新中...)