1 class Solution {

 2 public:

 3     //相当于图来存来处理

 4     vector<vector<int>> tree;

 5     vector<int> res;

 6     vector<int> subc;

 7     int n;

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

 9         //tree.rvec(N,vector<int>(N));

10         //init

11         n = N;

12         tree.resize(N);            //初始化的函数

13         res.assign(N,0);

14         subc.assign(N,0);

15         //build tree

16         for(auto e : edges){               //遍历的技巧

17             tree[e[0]].push_back(e[1]);

18             tree[e[1]].push_back(e[0]);

19         }

20         set<int> visited1;

21         set<int> visited2;

22         DFS_POST(0,visited1);         //初始root任何值都行

23         DFS_PRE(0,visited2);

24         return res;

25

26     }

27     void DFS_POST(int root,set<int> &visited){            //传引用保存修改值

28         visited.insert(root);

29         for(auto i : tree[root]){

30             if(visited.find(i) == visited.end() ){

31                 DFS_POST(i,visited);

32                 subc[root] += subc[i];

33                 res[root] += res[i] + subc[i];

34             }

35         }

36         subc[root]++;  //加上自身节点

37     }

38     void DFS_PRE(int root,set<int> &visited){

39         visited.insert(root);

40         for(auto i : tree[root]){

41             if(visited.find(i) == visited.end()){

42                 res[i] = res[root] - subc[i] + n - subc[i];         //算法核心

43                 DFS_PRE(i,visited);

44             }

45         }

46     }

47

48 };