mickypinata

LEET-T0834: Sum of Distances in Tree

Dec 1st, 2021
745
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3.  
  4. const int N = 3e4 + 5;
  5.  
  6. vector<vector<int>> edges;
  7. vector<int> adj[N];
  8. int subtree[N], dp[N], par[N];
  9.  
  10. void DFS(int u, int p){
  11.     par[u] = p;
  12.     subtree[u] = 1;
  13.     dp[u] = 0;
  14.     for(int v : adj[u]){
  15.         if(v == p){
  16.             continue;
  17.         }
  18.         DFS(v, u);
  19.         subtree[u] += subtree[v];
  20.         dp[u] += dp[v] + subtree[v];
  21.     }
  22. }
  23.  
  24. void DFSRev(vector<int> &ans, int nVertex, int u, int p){
  25.     if(p == -1){
  26.         ans[u] = dp[u];
  27.     } else {
  28.         ans[u] = nVertex + ans[p] - 2 * subtree[u];
  29.     }
  30.     for(int v : adj[u]){
  31.         if(v == p){
  32.             continue;
  33.         }
  34.         DFSRev(ans, nVertex, v, u);
  35.     }
  36. }
  37.  
  38. vector<int> sumOfDistancesInTree(int nVertex, vector<vector<int>> &edges){
  39.     for(vector<int> &p: edges){
  40.         int u = p[0];
  41.         int v = p[1];
  42.         adj[u].push_back(v);
  43.         adj[v].push_back(u);
  44.     }
  45.     DFS(0, -1);
  46.     vector<int> ans(nVertex, 0);
  47.     DFSRev(ans, nVertex, 0, -1);
  48.     return ans;
  49. }
  50.  
  51. int main(){
  52.  
  53.     int nVertex;
  54.     scanf("%d", &nVertex);
  55.     for(int i = 1; i < nVertex; ++i){
  56.         int u, v;
  57.         scanf("%d%d", &u, &v);
  58.         edges.push_back(vector<int>({u, v}));
  59.     }
  60.     for(int x : sumOfDistancesInTree(nVertex, edges)){
  61.         cout << x << ' ';
  62.     }
  63.  
  64.     return 0;
  65. }
  66.  
RAW Paste Data