leetcode834 Sum of Distances in Tree

 1 class Solution
 2 {
 3 public:
 4     void dfs(int root, int p, vector<vector<int>>& G, vector<int>& cnt, vector<int>& res)
 5     {
 6         for (auto it: G[root])
 7         {
 8             if (it == p) continue;
 9             dfs(it, root, G, cnt, res);
10             cnt[root] += cnt[it];
11             res[root] += res[it] + cnt[it];
12         }
13         cnt[root]++;
14     }
15     void dfs2(int root, int p, vector<vector<int>>& G, vector<int>& cnt, vector<int>& res, int N)
16     {
17         for (auto it: G[root])
18         {
19             if (it == p) continue;
20             res[it] = res[root] - cnt[it] + N - cnt[it];
21             dfs2(it, root, G, cnt, res, N);
22         }
23     }
24     vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges)
25     {
26         vector<vector<int>> G(N, vector<int>());
27         for (auto it: edges)
28         {
29             int a = it[0], b = it[1];
30             G[a].push_back(b); G[b].push_back(a);
31         }
32         vector<int> cnt(N, 0), res(N, 0);
33         dfs(0, -1, G, cnt, res);
34         dfs2(0, -1, G, cnt, res, N);
35         return res;
36     }
37 }