fork download
copy 
  1. #include<bits/stdc++.h>
  2. #define f(i,a,b) for(int i=a;i<b;i++)
  3. #define fe(i,a,b) for(int i=a;i<=b;i++)
  4. #define feh(a,b) for(auto a:b)
  5. #define pb push_back
  6. #define vi vector<int> 
  7. #define fast_io             ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
  8. using namespace std;
  9.  
  10. const int MAXN = 1e5 + 1;
  11.  
  12. struct Node
  13. {
  14.     vi ar;
  15. };
  16.  
  17. Node st[8*MAXN];
  18. vi adj[MAXN];
  19. int col[MAXN];
  20. int fa[2*MAXN];
  21. int s[MAXN];
  22. int e[MAXN];
  23. int N,Q,k,root;
  24. int ql,qr,ind,timer;
  25.  
  26. void merge(vi& A, vi& B, vi& res)
  27. {
  28.     int i=0,j=0;
  29.     while(i < A.size() && j < B.size())
  30.     {
  31.         if (A[i] < B[j]) 
  32.         {
  33.             res.pb(A[i]);
  34.             i++;
  35.         }
  36.         else if (A[i] == B[j])
  37.         {
  38.             if(res.empty() || res[res.size()-1] != A[i])
  39.                 res.pb(A[i]);
  40.  
  41.             i++;
  42.             j++;
  43.         }
  44.         else
  45.         {
  46.             res.pb(B[j]);
  47.             j++;
  48.         }
  49.     }
  50.  
  51.     f(k,i,A.size())
  52.     {
  53.         res.pb(A[k]); 
  54.     }
  55.  
  56.     f(k,j,B.size())
  57.     {
  58.         res.pb(B[k]); 
  59.     }
  60. }
  61.  
  62. void cst(int l, int r, int ind)
  63. {
  64.     if (l == r)
  65.     {
  66.         st[ind].ar.pb(col[fa[l]]);
  67.         // cout << "ind "<<ind<<" ===";
  68.  
  69.         // feh(val,st[ind].ar)
  70.         //     cout<<val<<" ";
  71.  
  72.         // cout<<"\n";
  73.          return;
  74.     }
  75.  
  76.     int mid = (l+r)/2;
  77.  
  78.     cst(l,mid,2*ind);
  79.     cst(mid+1,r,2*ind+1);
  80.  
  81.     merge(st[2*ind].ar,st[2*ind+1].ar,st[ind].ar);
  82.  
  83.     // cout << "ind "<<ind<<" ===";
  84.  
  85.     // feh(val,st[ind].ar)
  86.     //     cout<<val<<" ";
  87.  
  88.     // cout<<"\n";
  89. }
  90.  
  91. void rmq(int l, int r, int ind, vector<int>& res)
  92. {
  93.     if(l > qr || r < ql)
  94.     {
  95.         return ;
  96.     }
  97.  
  98.     if(l >= ql && r <= qr)
  99.     {
  100.         res = st[ind].ar;
  101.         return;
  102.     }
  103.  
  104.     int mid = (l+r)/2;
  105.  
  106.     vector<int> L,R;
  107.     rmq(l,mid,2*ind,L);
  108.     rmq(mid+1,r,2*ind+1,R); 
  109.     merge(L, R, res);
  110. }
  111.  
  112. void dfs(int v, int par)
  113. {
  114.     timer++;
  115.     fa[timer] = v;
  116.     s[v] = timer;
  117.  
  118.     feh(ver,adj[v])
  119.     {
  120.         if (ver != par)
  121.         {
  122.             dfs(ver,v);
  123.         }
  124.     }
  125.  
  126.     timer++;
  127.     fa[timer] = v;
  128.     e[v] = timer;
  129. }
  130.  
  131. int main()
  132. {
  133.     fast_io
  134.     int u,v;
  135.     cin >> N >> Q >> root;
  136.     timer = 0;
  137.     fe(i,1,N-1)
  138.     {
  139.         cin >> u >> v;
  140.         adj[u].pb(v);
  141.         adj[v].pb(u); 
  142.     }
  143.  
  144.     fe(i,1,N)
  145.     {
  146.         cin >> col[i];
  147.     }
  148.  
  149.     dfs(root, -1);
  150.  
  151.     cst(1,2*N,1);
  152.  
  153.     // fe(i,1,2*N)
  154.     //    cout <<fa[i]<<" ";
  155.  
  156.     // cout<<"\n"; 
  157.  
  158.     // fe(i,1,N)
  159.     //    cout<<s[i]<<" ";
  160.  
  161.     // cout<<"\n";
  162.  
  163.     // fe(i,1,N)
  164.     //    cout << e[i]<<" ";
  165.  
  166.     // cout<<"\n";
  167.  
  168.     while(Q--)
  169.     {
  170.         vector<int> res;
  171.         cin >> k;
  172.         ql = s[k];
  173.         qr = e[k];
  174.         rmq(1,2*N,1,res);
  175.         cout << res.size() <<"\n";
  176.         // for(int val:res)
  177.         //    cout<<val<<" ";
  178.         //    cout<<"\n";
  179.         //cout << res.size() <<"\n";
  180.     }
  181.  
  182.     return 0;
  183. }
  184.  
Success #stdin #stdout 0.01s 24572KB
comments (?)
stdin
copy 
4 2 1
1 2
2 4
2 3
10
20
20
30
1
2
stdout
copy 
3
2
https://ideone.com/XiHttO
Coloring tree hackerrank
https://w...content-available-to-author-only...k.com/challenges/coloring-tree/problem
language:
C++14 (gcc 8.3)
created:
5 years ago
visibility:
public

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!