题目描述

给定一个无向、连通的树。树中有 N 个标记为 0...N-1 的节点以及 N-1 条边 。

第 i 条边连接节点 edges[i][0] 和 edges[i][1]。

返回一个表示节点 i 与其他所有节点距离之和的列表 ans。

样例

输入：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

算法

(DFS) $O(n)$

1. 这里需要两次 DFS。假设 0 号结点为根结点，将整棵树变为有根树。
2. 第一次 DFS 求出每个点子树的大小（包含自身），以及该结点到子树所有结点的距离之和。
3. 第二次 DFS 需要累计每个结点从父亲获得的其他结点的距离之和。
4. 具体地，第一次 DFS 时，假设当前结点是 $x$，某个儿子是 $y$，则$num[x] = num[y] + num[y]$，$ans[x] = ans[x] + num[y]$，最后 $num[x] = num[x] + 1$。
5. 第二次 DFS 时，仍然假设当前结点为 $x$，某个儿子是 $y$，并且假设 $ans[x]$ 就是其他所有结点到当前结点的距离之和，接下来更新 $ans[y]$。更新之前 $ans[y]$ 是 $y$ 及其子结点到达它的距离之和，需要补全经过 $x$ 到 $y$ 的结点：$ans[y] = ans[y] + (ans[x] - (ans[y] + num[y]) + (N - num[y]))$，即用 $ans[x]$ 除去经过 $y$ 到达 $x$ 的结点，然后将剩余的结点整体向 $y$ 移动一步，移动所带来的代价就是剩余的结点数 $N - num[y]$，就像第一次 DFS 更新时，每次所有儿子结点向当前结点移动一步所带来的更新就是儿子结点的个数。

时间复杂度

• 每个结点仅遍历两次，故时间复杂度为 $O(n)$。

C++ 代码

class Solution {
public:

    void dfs1(int x, int fa, const vector<vector<int>>& graph,
              vector<int> &ans, vector<int> &num) {

        for (auto y : graph[x])
            if (y != fa) {
                dfs1(y, x, graph, ans, num);
                num[x] += num[y];
                ans[x] += ans[y] + num[y];
            }

        num[x]++;
    }

    void dfs2(int x, int fa, const vector<vector<int>>& graph,
              vector<int> &ans, const vector<int> &num, int N) {

        for (auto y : graph[x])
            if (y != fa) {
                ans[y] += ans[x] - ans[y] - num[y] + N - num[y];
                dfs2(y, x, graph, ans, num, N);
            }

    }

    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {

        vector<vector<int>> graph(N);

        for (int i = 0; i < N - 1; i++) {
            graph[edges[i][0]].push_back(edges[i][1]);
            graph[edges[i][1]].push_back(edges[i][0]);
        }

        vector<int> num(N, 0), ans(N, 0);
        dfs1(0, -1, graph, ans, num);
        dfs2(0, -1, graph, ans, num, N);

        return ans;
    }
};