LeetCode 树中距离之和

给定一个无向、连通的树。树中有 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 $思路分析：$这道题题意是比较简单的，蛋式数据测试量比较大。它就是让我们确定以每个点作为根节点，此时根节点到各个点的距离和。 深度优先搜索法（未通过）

class Solution { public: vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) { vector<vector<int>> myTree(N);//myTree[i]用于记录i能到达的各个点 for (const auto &edge : edges) {//将edge信息转换myTree myTree[edge[0]].push_back(edge[1]); myTree[edge[1]].push_back(edge[0]); } vector<int> resVec(N, 1);//存储结果，resVec[i]记录以i作为根节点，到达各个点的距离 for (int i = 0; i < N; ++i) { myDFS(-1, i, 1, resVec[i], myTree); } return resVec; } //深度优先搜索以nowIndex为根（beforeIndex是nowIndex的父节点）到达它子节点的距离之和 void myDFS(int beforeIndex, int nowIndex, int depth, int &sumRes, vector<vector<int>> &myTree) { sumRes += (myTree[nowIndex].size() - 1) * depth;//到达子节点的距离 for (auto &nextIndex : myTree[nowIndex]) { if (nextIndex != beforeIndex) {//由子节点继续搜索 myDFS(nowIndex, nextIndex, depth + 1, sumRes, myTree); } } } };

经过翻阅前辈的解题思路，发现一种NB的算法。 对于上面这个图，我们将xy这条边分隔开，我们定义以左边x为根的所有节点距离和是x@X，右边以y为根的所有节点距离和是y@Y，而x到右边y的所有点的距离和定义为x@Y，那么此时我们以x为根的结果就应该是x@X+x@Y，而x@Y=y@Y+count(y)（其中count(y)表示y的节点个数），那么

res(x)=x@X+y@Y+count(y)res(x)=x@X+y@Y+count(y)

同理，此时如果我们将y作为根的话，也可以得到一个类似的结论

res(y)=x@X+y@Y+count(x)res(y)=x@X+y@Y+count(x)

那么，我们就很容易知道

$res(x)-res(y)=count(y)-count(x)res(x)-res(y)=count(y)-count(x)$

这是一个非常有用的结论。

接着的问题就是怎么计算节点的数量和节点到它所有孩子的距离。节点数量很好计算就是将所有孩子的包含的节点数加起来然后再加上自身就行了，我们定义count(node)表示节点数，那么

count(node) = count(childs) + 1

接着考虑当前节点到其所有子节点的距离和。这个也非常简单，实际上和上面是一样的，我们只需要将所有孩子到其所有子节点的和加起来，然后再加上当前节点的子节点个数就行了，我们定义sum(node)表示距离和，那么

sum(node)=sum(childs) + count(childs)sum(node)= sum(childs) + count(childs)

那么我们只需要通过一次深度遍历就可以求得所有节点的count，但是所有的sum我们是无法求得的，因为我们只可以求和当前节点到其所有子节点的和，这个时候我们就需要联系这个节点到这个节点的子节点以外的所有节点的距离，这就要用到我们前面那个结论。

我们可以知道 $res(child)-res(parent)=count(parent)-count(child)res(child)-res(parent)=count(parent)-count(child)$

我们知道count(parent)=N - count(child)，所以 $res(parent)-res(child)=2\ast count(child)-Nres(parent)-res(child)=2\ast count(child)-N$

class Solution { public: vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) { vector<vector<int>> myTree(N);//myTree[i]用于记录i能到达的各个点 for (const auto &edge : edges) {//将edge信息转换myTree myTree[edge[0]].push_back(edge[1]); myTree[edge[1]].push_back(edge[0]); } //count[i]记录i的子节点个数 vector<int> count(N, 1), resVec(N, 0); DFS(myTree, count, resVec, 0, -1); dfs(myTree, resVec, count, N, 0, -1); return resVec; } //计算nowIndex为根（beforeIndex为nowIndex的父节点）的子节点个数 void DFS(vector<vector<int>> &myTree, vector<int> &count, vector<int> &resVec, int nowIndex, int beforeIndex) { for (const auto &nextIndex : myTree[nowIndex]) { if (nextIndex != beforeIndex){ DFS(myTree, count, resVec, nextIndex, nowIndex); count[nowIndex] += count[nextIndex];//统计nowIndex的子节点个数 resVec[nowIndex] += resVec[nextIndex] + count[nextIndex]; } } } //计算nowIndex为根（beforeIndex为nowIndex的父节点）的结果 void dfs(vector<vector<int>> &myTree, vector<int> &resVec, vector<int> &count, int &N, int nowIndex, int beforeIndex) { for (const auto &nextIndex : myTree[nowIndex]) { if (nextIndex != beforeIndex){ resVec[nextIndex] = resVec[nowIndex] - 2 * count[nextIndex] + N;//套最后一条公式（注意nowIndex是nextIndex的parent） dfs(myTree, resVec, count, N, nextIndex, nowIndex); } } } };