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.

class Solution {
public:
    vector<vector<int>> gh;
    vector<int> dist; //i点的所有子结点到i的距离之和
    vector<int> size; //以i点为根的树的结点数量和(包括i结点本身)
    vector<int> res;
    vector<int> sumOfDistancesInTree(int n, vector<vector<int>>& edges) {
        gh.resize(n); dist.resize(n); size.resize(n); res.resize(n);
        for (const auto& v : edges) {
            gh[v[0]].emplace_back(v[1]);
            gh[v[1]].emplace_back(v[0]);
        }
        dfs1(0, -1);
        res[0] = dist[0];
        dfs2(n, 0, -1);
        return res;
    }
    //求出以0为根的距离之和且求出了分别以i为根的子树结点数量和
    void dfs1(int cur, int par) {
        dist[cur] = 0;
        size[cur] = 1;
        for (auto i : gh[cur]) {
            if (i == par) continue;
            dfs1(i, cur);
            dist[cur] += dist[i] + size[i];
            size[cur] += size[i];
        }
    }
    //从0开始往下递归推出子树。再子树推算出子子树直到所有。
    void dfs2(int n, int cur, int par) {
        for (auto i : gh[cur]) {
            if (i == par) continue;
            res[i] = res[cur] + (n - size[i]) - size[i];
            dfs2(n, i, cur);
        }
    }
};