[leetcode]834. Sum of Distances in Tree

链接：https://leetcode.com/problems/sum-of-distances-in-tree/description/

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.

Note: 1 <= N <= 10000

 class Solution {
 public:
      //相当于图来存来处理
      vector<vector<int>> tree;
      vector<int> res;
      vector<int> subc;
      int n;
      vector<int> sumOfDistancesInTree(int N, vector<vector<int>>& edges) {
          //tree.rvec(N,vector<int>(N));
         //init
         n = N;
         tree.resize(N);            //初始化的函数
         res.assign(N,0);
         subc.assign(N,0);
         //build tree
         for(auto e : edges){               //遍历的技巧
             tree[e[0]].push_back(e[1]);
             tree[e[1]].push_back(e[0]);
         }
         set<int> visited1;
         set<int> visited2;
         DFS_POST(0,visited1);         //初始root任何值都行
         DFS_PRE(0,visited2);
         return res;
         
     }
     void DFS_POST(int root,set<int> &visited){            //传引用保存修改值
         visited.insert(root);
         for(auto i : tree[root]){
             if(visited.find(i) == visited.end() ){
                 DFS_POST(i,visited);
                 subc[root] += subc[i];
                 res[root] += res[i] + subc[i];    
             }
         }
         subc[root]++;  //加上自身节点
     }
     void DFS_PRE(int root,set<int> &visited){
         visited.insert(root);
         for(auto i : tree[root]){
             if(visited.find(i) == visited.end()){
                 res[i] = res[root] - subc[i] + n - subc[i];         //算法核心
                 DFS_PRE(i,visited);
             }
         }
     }
     
 };