fork(13) download
copy 
  1. #include<bits/stdc++.h>
  2. #define ll long long
  3. #define endl '\n'
  4. #define watch(x) cout<<(#x)<<" = "<<x<<endl
  5. #define arr_watch(arr, n) for(int i=0;i<n;++i) cout<<arr[i]<<" \n"[i==n-1]
  6. using namespace std;
  7. const int maxn = 1e5;
  8. vector<int> graph[maxn + 5];
  9. int order[maxn + 5], A[maxn + 5];
  10. list<int> colors[maxn + 5];
  11. int the_end;
  12.  
  13. void dfs(int u, int p){
  14.     colors[A[u]].push_back(u);
  15.     for (int child: graph[u]){
  16.         if(child == p) continue;
  17.         dfs(child, u);
  18.     }
  19.     order[u] = the_end++;
  20. }
  21.  
  22. //LM10_Piyush...
  23. int main(){
  24.     ios_base::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
  25.     #ifndef ONLINE_JUDGE
  26.  
  27.     // freopen("D:/TestFiles/input.txt","r",stdin);
  28.  
  29.     #endif
  30.     /*********************** CODE IS HERE *****************************************/
  31.  
  32.     //number of nodes...
  33.     int n; cin >> n;
  34.     for(int i = 1, x, y; i < n; ++i){
  35.         cin >> x >> y;
  36.         graph[x].push_back(y);
  37.         graph[y].push_back(x);
  38.     }
  39.     //colors input...
  40.     for (int i = 1; i <= n; ++i)
  41.         cin >> A[i];    
  42.  
  43.     the_end = 1;
  44.     dfs(1, 0);        //here we go....
  45.  
  46.     int ans = 0;
  47.     int q; cin >> q;
  48.     while(q--){
  49.         int x; cin >> x;
  50.         int col = A[x];   //color of that x node...
  51.         auto it1 = colors[col].begin();
  52.         auto it2 = it1;
  53.         bool flag = false;
  54.  
  55.         //iterate over subtree with same color(end_time of those nodes <= xth node)
  56.         while(it1 != colors[col].end()){
  57.             if(order[*it1] > order[x] and flag)
  58.                 break;
  59.             it2++;
  60.             if(*it1 == x)
  61.                 flag = true;
  62.  
  63.             if(order[*it1] <= order[x] and flag){
  64.                 ans++;
  65.                 colors[col].erase(it1);  //found that node, then remove it
  66.             }
  67.             it1 = it2;
  68.         }
  69.         cout << ans << endl;
  70.  
  71.     }
  72.  
  73.  
  74.     return 0;
  75. }
  76.  
Success #stdin #stdout 0s 8116KB
comments (?)
stdin
copy 
5
3 1
3 2
2 4
2 5
1 1 2 2 1
4
2
4
5
1
stdout
copy 
2
3
3
4
https://ideone.com/dL7Qhm
Hakerarth ques Hackwithinfy
language:
C++14 (gcc 8.3)
created:
4 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!