834. 树中距离之和 dfs

给定一个无向、连通的树。树中有 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

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/sum-of-distances-in-tree
著作权归领扣网络所有。商业转载请联系官方授权，非商业转载请注明出处。

ans[i] : i节点到其他所有节点的距离
sum[i] : 以i节点为root的子树里的所有节点到i的距离
cnt[i] : 以i节点为root的子树的大小，包括i自己
将问题拆分：对于两个相邻节点A和B，将树拆分为两个子树，根节点分别为A和B，ans(A) = sum(A) + sum(B) + cnt(B);
*如何理解cnt(b)？因为A、B相邻，已经求出B到其子树里所有节点的距离，对A来说，对每个节点都+1，就变成了A到B子树里所有节点的距离。
同理，ans(B) = sum(B) + sum(A) + cnt(A);
由此我们得到：
ans(A) = sum(A) + sum(B) + cnt(B);
ans(B) = sum(B) + sum(A) + cnt(A);
ans(A) = ans(B) - cnt(A) + cnt(B) = ans(B) - cnt(A) + (N - cnt(A));
因此，对于根节点root的任意子节点child，ans(child) = ans(root) - cnt(child) + N - cnt(child);
得到root的答案就可以DFS递归得到其他所有节点的答案。（这里需要一个DFS）

class Solution {
public:
    vector <int> ans;
    vector <int> cnt;
    vector <vector <int>> G;
    int n;

    void dfs(int now, int fa) {
        for (int i = 0; i < G[now].size(); i++) {
            int v = G[now][i];
            if (v != fa) {
                dfs(v, now);
                cnt[now] += cnt[v];
                ans[now] += ans[v] + cnt[v];
            }
        }
    }

    void dfs2(int now, int fa) {
        for (int i =  0; i < G[now].size(); i++) {
            int v = G[now][i];
            if (v != fa) {
                ans[v] = ans[now] - cnt[v] + (n - cnt[v]);
                dfs2(v, now);
            }
        }
    }

    vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
        if (N == 0 || edges.empty()) {
            return {0};
        }

        n = N;
        G.resize(N);
        ans.resize(N);
        cnt.resize(N);

        for (int i = 0; i < n; i++) {
            cnt[i] = 1;
        }
        for (auto edge : edges) {
            int u = edge[0];
            int v = edge[1];
            G[u].push_back(v);
            G[v].push_back(u);
        }

        //第一次dfs求得节点0的正确答案
        //第一次dfs后的ans[i]表示节点i到其所有子节点的距离之和
        dfs(0, -1);

        //第二次dfs通过节点0的答案求得其他节点的正确答案
        //第二次dfs后的ans[i]表示节点i到其所有 节点的距离之和
        dfs2(0 , -1);

        return ans;
    }
};