Advertisement
humbertoyusta

Untitled

Mar 15th, 2022
253
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 4.95 KB | None | 0 0
  1. #include<bits/stdc++.h>
  2. using namespace std;
  3. // Macros:
  4. #define f first
  5. #define s second
  6. typedef long long ll;
  7. typedef long double ld;
  8. // debugging macros
  9. #define db(x) cerr << #x << ": " << (x) << '\n';
  10. #define db_v(x) cerr << #x << ": ["; for( auto i : (x) ) cerr << i << ", "; cerr << "]\n";
  11. #define db_p(x) cerr << #x << ": " << (x).first << " " << (x).second << '\n';
  12. // random
  13. mt19937_64 rnd(time(0));
  14. unsigned long long rng (unsigned long long a, unsigned long long b){
  15.     return uniform_int_distribution<unsigned long long>(a,b)(rnd);
  16. }
  17.  
  18. unsigned long long allK;
  19. vector<unsigned long long> rand_val;
  20. vector<int> color;
  21.  
  22. typedef vector<vector<int>> graph;
  23. struct centroid{
  24.     graph G;
  25.     vector<int> subt, boss;
  26.     vector<unsigned long long> hash_;
  27.     vector<bool> mk;
  28.     int size_subt, n, root;
  29.     set<unsigned long long> S;
  30.     vector<int> good_path, answer;
  31.  
  32.     centroid(graph _G) {
  33.         G = _G; n = G.size();
  34.         subt.resize(n); mk.resize(n); boss.resize(n), hash_.resize(n);
  35.         good_path.resize(n), answer.resize(n);
  36.         for(int i=1; i<n; i++)
  37.             if( !mk[i] ){
  38.                 build_centroid(i,0);
  39.             }
  40.     }
  41.  
  42.     void dfs_prec(int u,int p){
  43.         subt[u] = 1;
  44.         for( auto v : G[u] )
  45.             if( !mk[v] && v != p ){
  46.                 dfs_prec(v,u);
  47.                 subt[u] += subt[v];
  48.             }
  49.     }
  50.  
  51.     /// specific code begins
  52.     void dfs_compute(int u,int p)
  53.     {
  54.         if (u != p)
  55.         {
  56.             /// compute the hash of the path from each centroid to each node
  57.             hash_[u] = hash_[p] + rand_val[color[u]];
  58.         }
  59.         /// insert the hashes in a set
  60.         S.insert(hash_[u]);
  61.  
  62.         for( auto v : G[u] )
  63.             if (!mk[v] && v != p)
  64.             {
  65.                 dfs_compute(v, u);
  66.             }
  67.     }
  68.  
  69.     void dfs_solve(int u,int p,int centroid)
  70.     {
  71.         /// compute the hash needed to complement the path from the centroid to
  72.         /// the current node u, subtracting this path and the centroid from 1 2 3 ... k
  73.         unsigned long long hash_needed = allK - hash_[u] - rand_val[color[centroid]];
  74.  
  75.         /// check if the hash_needed is in the set, to mark all the nodes from u to centroid as solved
  76.         /// for each path (u, v) in subt(centroid) that is a permutation, all its nodes will be marked
  77.         /// since this part of the code will be true for (u, centroid) and for (v, centroid)
  78.         auto it = S.lower_bound(hash_needed);
  79.         if( it != S.end() && (*it) == hash_needed )
  80.         {
  81.             good_path[u] ++;
  82.         }
  83.  
  84.         for( auto v : G[u] )
  85.             if (!mk[v] && v != p)
  86.             {
  87.                 dfs_solve(v, u, centroid);
  88.             }
  89.     }
  90.  
  91.     void dfs_accumulate(int u,int p)
  92.     {
  93.         for (auto v : G[u])
  94.         {
  95.             if (!mk[v] && v != p)
  96.             {
  97.                 dfs_accumulate(v, u);
  98.                 /// make a bottom up accumulation to mark all the nodes from u to centroid as solved
  99.                 good_path[u] += good_path[v];
  100.                 good_path[v] = 0;
  101.             }
  102.         }
  103.         /// if a node is in any good path, then its answer is YES
  104.         if (good_path[u] > 0)
  105.             answer[u] = 1;
  106.         if (u == p)
  107.             good_path[u] = 0;
  108.     }
  109.     /// specific code ends
  110.  
  111.     int find_centroid(int u,int p){
  112.         for( auto v : G[u] )
  113.             if( !mk[v] && subt[v] > size_subt / 2 && v != p )
  114.                 return find_centroid(v,u);
  115.         return u;
  116.     }
  117.  
  118.     void build_centroid(int nod,int p){
  119.         dfs_prec(nod,nod);
  120.         size_subt = subt[nod];
  121.         int centroid = find_centroid(nod,nod);
  122.  
  123.         /// specific code
  124.         S.clear();
  125.         hash_[centroid] = boss[centroid] = 0;
  126.  
  127.         dfs_compute(centroid, centroid);
  128.  
  129.         dfs_solve(centroid, centroid, centroid);
  130.  
  131.         dfs_accumulate(centroid, centroid);
  132.         /// /specific code
  133.  
  134.         mk[centroid] = 1;
  135.         for( auto u : G[centroid] ){
  136.             if( !mk[u] ){
  137.                 build_centroid(u,centroid);
  138.             }
  139.         }
  140.     }
  141. };
  142.  
  143. int main(){
  144.     ios_base::sync_with_stdio(0); cin.tie(0);
  145.     ///cout.setf(ios::fixed); cout.precision(10);
  146.  
  147.     //freopen("a.in","r",stdin);
  148.     //freopen("a.in","w",stdout);
  149.  
  150.     int n, k, q;
  151.     cin >> n >> k >> q;
  152.  
  153.     rand_val.resize(k+1);
  154.     allK = 0;
  155.     for(int i=1; i<=k; i++)
  156.     {
  157.         rand_val[i] = rng(0,1ll<<62);
  158.         allK += rand_val[i];
  159.     }
  160.  
  161.     color.resize(n+1);
  162.     for(int i=1; i<=n; i++)
  163.     {
  164.         cin >> color[i];
  165.     }
  166.  
  167.     graph G(n+1);
  168.     for(int i=1; i<n; i++)
  169.     {
  170.         int u, v;
  171.         cin >> u >> v;
  172.         G[u].push_back(v);
  173.         G[v].push_back(u);
  174.     }
  175.  
  176.     centroid cTree = centroid(G);
  177.  
  178.     for(int i=1; i<=q; i++)
  179.     {
  180.         int u;
  181.         cin >> u;
  182.         if( cTree.answer[u] )
  183.             cout << "YES\n";
  184.         else
  185.             cout << "NO\n";
  186.     }
  187.  
  188.     return 0;
  189. }
  190.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement