[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 题目讲解汇总(持续更新中...)](https://www.cnblogs.com/grandyang/p/4606334.html)