Advertisement
Ankit_132

D

Sep 25th, 2023
276
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 2.25 KB | None | 0 0
  1.  
  2. #include <bits/stdc++.h>
  3.  
  4. using namespace std;
  5.  
  6. #define ll     long long
  7. #define _test   int _TEST; cin>>_TEST; while(_TEST--)
  8.  
  9. int main()
  10. {
  11.     _test
  12.     {
  13.         int n;
  14.         cin>>n;
  15.        
  16.         vector<ll int> a(n);
  17.         for(auto &e: a) cin>>e;
  18.        
  19.         vector<vector<int>> tree(n);
  20.         for(int i=0; i<n-1; i++)
  21.         {
  22.             int u, v;
  23.             cin>>u>>v;
  24.            
  25.             u--, v--;
  26.            
  27.             tree[u].push_back(v);
  28.             tree[v].push_back(u);
  29.         }
  30.        
  31.         vector<ll int> cnt(n);
  32.        
  33.         function<void(int, int)> _DFS = [&](int u, int p)
  34.         {
  35.             for(auto v: tree[u])
  36.             {
  37.                 if(v == p)      continue;
  38.                
  39.                 _DFS(v, u);
  40.                 cnt[u] += cnt[v];
  41.             }
  42.            
  43.             cnt[u]++;
  44.         };
  45.        
  46.         _DFS(0, -1);
  47.        
  48.         vector<ll int> val(n);
  49.        
  50.         function<void(int, int)> DFS = [&](int u, int p)
  51.         {
  52.             for(auto v: tree[u])
  53.             {
  54.                 if(v == p)      continue;
  55.                
  56.                 DFS(v, u);
  57.                
  58.                 val[u] += (a[u]^a[v])*1ll*cnt[v] + val[v];
  59.             }
  60.         };
  61.        
  62.         DFS(0, -1);
  63.        
  64.         vector<ll int> ans(n);
  65.        
  66.         function<void(int, int)> DFS2 = [&](int u, int p)
  67.         {
  68.             ans[u] = val[u];
  69.            
  70.             for(auto v: tree[u])
  71.             {
  72.                 if(v == p)      continue;
  73.                
  74.                 val[u] -= (a[u]^a[v])*1ll*cnt[v];
  75.                 val[u] -= val[v];
  76.                 cnt[u] -= cnt[v];
  77.                 cnt[v] += cnt[u];
  78.                 val[v] += (a[u]^a[v])*1ll*cnt[u];
  79.                 val[v] += val[u];
  80.                
  81.                 DFS2(v, u);
  82.                
  83.                 val[v] -= val[u];
  84.                 val[v] -= (a[u]^a[v])*1ll*cnt[u];
  85.                 cnt[v] -= cnt[u];
  86.                 cnt[u] += cnt[v];
  87.                 val[u] += val[v];
  88.                 val[u] += (a[u]^a[v])*1ll*cnt[v];
  89.             }
  90.         };
  91.        
  92.         DFS2(0, -1);
  93.        
  94.         for(auto e: ans)        cout<<e<<" ";
  95.         cout<<"\n";
  96.     }
  97. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement