fork download
copy 
  1. #include <iostream>
  2. #include <algorithm>
  3. #include <cstdio>
  4. #include <vector>
  5.  
  6. using namespace std;
  7.  
  8. int N, LK, RK, LC, RC;
  9. char A[200002];
  10. vector<int> adj[200001];
  11. bool bad[200001];
  12. int sz[200001];
  13. vector<vector<pair<int, int>>> V;
  14. vector<pair<int, int>> all;
  15. int bit[200006];
  16. long long ans;
  17.  
  18. void add(int x, int v)
  19. {
  20.     for(x++; x<=200005; x+=x&-x)
  21.         bit[x]+=v;
  22. }
  23.  
  24. int sum(int x)
  25. {
  26.     int ret=0;
  27.     for(x++; x>0; x-=x&-x)
  28.         ret+=bit[x];
  29.     return ret;
  30. }
  31.  
  32. void dfs(int u, int p)
  33. {
  34.     sz[u]=1;
  35.     for(auto& v: adj[u]) if(v!=p && !bad[v])
  36.     {
  37.         dfs(v, u);
  38.         sz[u]+=sz[v];
  39.     }
  40. }
  41.  
  42. void dfs2(int u, int p, int K, int C)
  43. {
  44.     if(A[u]=='K')
  45.         K++;
  46.     else
  47.         C++;
  48.     V.back().push_back(make_pair(K, C));
  49.     all.push_back(make_pair(K, C));
  50.     for(auto& v: adj[u]) if(v!=p && !bad[v])
  51.         dfs2(v, u, K, C);
  52. }
  53.  
  54. long long work(vector<pair<int, int>>& v, int R0, int R1)
  55. {
  56.     if(R0<0 || R1<0)
  57.         return 0;
  58.     long long ret=0;
  59.     int j=0;
  60.     for(int i=v.size()-1; i>=0; i--)
  61.     {
  62.         while(j<(int)v.size() && v[i].first+v[j].first<=R0)
  63.         {
  64.             add(v[j].second, 1);
  65.             j++;
  66.         }
  67.         ret+=sum(R1-v[i].second);
  68.         if(j>=i && v[i].first*2<=R0 && v[i].second*2<=R1)
  69.             ret--;
  70.     }
  71.     for(int i=0; i<j; i++)
  72.         add(v[i].second, -1);
  73.     return ret/2;
  74. }
  75.  
  76. long long work(vector<pair<int, int>>& v)
  77. {
  78.     sort(v.begin(), v.end());
  79.     return work(v, RK, RC)-work(v, LK-1, RC)-work(v, RK, LC-1)+work(v, LK-1, LC-1);
  80. }
  81.  
  82. void solve(int u)
  83. {
  84.     dfs(u, -1);
  85.     while(1)
  86.     {
  87.         int idx=-1;
  88.         for(auto& v: adj[u]) if(!bad[v] && (idx==-1 || sz[v]>sz[idx]))
  89.             idx=v;
  90.         if(sz[idx]*2<=sz[u])
  91.             break;
  92.         sz[u]-=sz[idx];
  93.         sz[idx]+=sz[u];
  94.         u=idx;
  95.     }
  96.     if((A[u]=='K' && LK<=1 && LC==0) || (A[u]=='C' && LC<=1 && LK==0))
  97.         ans++;
  98.     if((A[u]=='K' && RK==0) || (A[u]=='C' && RC==0))
  99.         goto skip;
  100.     if(A[u]=='K')
  101.         LK--, RK--;
  102.     else
  103.         LC--, RC--;
  104.     V.clear();
  105.     all.clear();
  106.     for(auto& v: adj[u]) if(!bad[v])
  107.     {
  108.         V.push_back(vector<pair<int, int>>());
  109.         dfs2(v, u, 0, 0);
  110.         for(auto& it: V.back())
  111.             if(LK<=it.first && it.first<=RK && LC<=it.second && it.second<=RC)
  112.                 ans++;
  113.     }
  114.     ans+=work(all);
  115.     for(auto& it: V)
  116.         ans-=work(it);
  117.     if(A[u]=='K')
  118.         LK++, RK++;
  119.     else
  120.         LC++, RC++;
  121.     skip:;
  122.     bad[u]=true;
  123.     for(auto& v: adj[u]) if(!bad[v])
  124.         solve(v);
  125. }
  126.  
  127. int dfs_bf(int u, int p, int K, int C)
  128. {
  129.     if(A[u]=='K')
  130.         K++;
  131.     else
  132.         C++;
  133.     int ret=0;
  134.     if(LK<=K && K<=RK && LC<=C && C<=RC)
  135.         ret++;
  136.     for(auto& v: adj[u]) if(v!=p)
  137.         ret+=dfs_bf(v, u, K, C);
  138.     return ret;
  139. }
  140.  
  141. long long solve_bf()
  142. {
  143.     long long ret=0;
  144.     for(int i=1; i<=N; i++)
  145.     {
  146.         ret+=dfs_bf(i, -1, 0, 0);
  147.         if(LK<=(A[i]=='K') && (A[i]=='K')<=RK && LC<=(A[i]=='C') && (A[i]=='C')<=RC)
  148.             ret++;
  149.     }
  150.     return ret/2;
  151. }
  152.  
  153. int main()
  154. {
  155.     scanf("%d%d%d%d%d", &N, &LK, &RK, &LC, &RC);
  156.     scanf("%s", A+1);
  157.     int a, b;
  158.     for(int i=1; i<N; i++)
  159.     {
  160.         scanf("%d%d", &a, &b);
  161.         adj[a].push_back(b);
  162.         adj[b].push_back(a);
  163.     }
  164.     solve(1);
  165.     printf("%lld\n", ans);
  166.     //printf("%lld\n", solve_bf());
  167.     return 0;
  168. }
  169.  
Success #stdin #stdout 0s 7448KB
comments (?)
stdin
copy 
Standard input is empty
stdout
copy 
0
https://ideone.com/L50nyo
language:
C++14 (gcc 8.3)
created:
9 years ago
visibility:
secret

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!