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 题目讲解汇总(持续更新中…)


转载请注明来源于 Grandyang 的博客 (grandyang.com),欢迎对文章中的引用来源进行考证,欢迎指出任何有错误或不够清晰的表达。可以在下面评论区评论,也可以邮件至 grandyang@qq.com

💰


微信打赏


Venmo 打赏

(欢迎加入博主的知识星球,博主将及时答疑解惑,并分享刷题经验与总结,试运营期间前五十位可享受半价优惠~)

×

Help us with donation