Guest User

Untitled

a guest
Feb 17th, 2020
407
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 4.54 KB | None | 0 0
  1. #include <bits/stdc++.h>
  2.  
  3. using namespace std;
  4.  
  5.  
  6. struct CentroidDecomposer {
  7.     int n;
  8.     vector<set<int>> G;
  9.     vector<int> Size;
  10.     // Depth and parent in the centroid tree
  11.     vector<int> CD, CP;
  12.     vector<vector<int>> dp, depth;
  13.     vector<bool> tag;
  14.     vector<int> all;
  15.     vector<int> down, up;
  16.     int k;
  17.  
  18.     CentroidDecomposer(int n, int k) :
  19.         n(n), k(k), G(n), Size(n), CD(n), CP(n), tag(n, false),
  20.         down(n, -1), up(n, -1),
  21.         depth(20, vector<int>(n)), dp(20, vector<int>(n)) {}
  22.  
  23.     void Tag(int node) {
  24.         tag[node] = true;
  25.     }
  26.  
  27.     void AddEdge(int a, int b) {
  28.         G[a].insert(b);
  29.         G[b].insert(a);
  30.     }
  31.  
  32.     bool Solve(int a, int b) {
  33.         int na = a, nb = b;
  34.         while (CD[na] > CD[nb])
  35.             na = CP[na];
  36.         while (CD[na] < CD[nb])
  37.             nb = CP[nb];
  38.         while (na != nb) {
  39.             na = CP[na];
  40.             nb = CP[nb];
  41.         }
  42.         int cen = na;
  43.         int da = dp[CD[cen]][a];
  44.         int db = dp[CD[cen]][b];
  45.         if (da + db <= k) {
  46.             return true;
  47.         }
  48.         return false;
  49.     }
  50.  
  51.     int rec_size, rec_centr;
  52.     void ComputeSizeAndCentroid(int node, int par) {
  53.         Size[node] = 1;
  54.         int max_size = 0;
  55.  
  56.         for(auto vec : G[node])
  57.             if(vec != par) {
  58.                 ComputeSizeAndCentroid(vec, node);
  59.                 Size[node] += Size[vec];
  60.                 max_size = max(max_size, Size[vec]);
  61.             }
  62.  
  63.         max_size = max(max_size, rec_size - Size[node]);
  64.         if(2 * max_size <= rec_size)
  65.             rec_centr = node;
  66.     }
  67.  
  68.     void DoUsefulDFS(int node, int par, int cd) {
  69.         Size[node] = 1;
  70.         all.push_back(node);
  71.         for(auto vec : G[node])
  72.             if(vec != par) {
  73.                 depth[cd][vec] = 1 + depth[cd][node];
  74.                 DoUsefulDFS(vec, node, cd);
  75.                 Size[node] += Size[vec];
  76.             }
  77.     }
  78.  
  79.     void DFS1(int cd, int node, int par) {
  80.         if (tag[node]) {
  81.             down[node] = node;
  82.         }
  83.  
  84.         for (auto vec : G[node]) {
  85.             if (vec == par) continue;
  86.             DFS1(cd, vec, node);
  87.             if (down[vec] != -1 && (down[node] == -1 ||
  88.                     depth[cd][down[node]] > depth[cd][down[vec]])) {
  89.                 down[node] = down[vec];
  90.             }
  91.         }
  92.     }
  93.  
  94.     multimap<int, int> DFS2(int cd, int node, int par) {
  95.         multimap<int, int> big;
  96.         big.insert({depth[cd][node], node});
  97.  
  98.         for (auto vec : G[node]) {
  99.             if (vec == par) continue;
  100.  
  101.             auto small = DFS2(cd, vec, node);
  102.             if (small.size() > big.size()) swap(small, big);
  103.             for (auto x : small) {
  104.                 big.insert(x);
  105.             }
  106.         }
  107.  
  108.         if (down[node] != -1) {
  109.             int dd = depth[cd][down[node]];
  110.             int dr = depth[cd][node];
  111.             // dd + x - 2 * dr <= k
  112.             // x <= k + 2 * dr - dd
  113.             int maxd = k + 2 * dr - dd;
  114.             int mind = dd + 1;
  115.             auto it = big.lower_bound(mind);
  116.             while (it != big.end() && it->first <= maxd) {
  117.                 assert(up[it->second] == -1);
  118.                 up[it->second] = down[node];
  119.                 it = big.erase(it);
  120.             }
  121.         }
  122.         return big;
  123.     }
  124.  
  125.     void BFS(int cd, int root) {
  126.  
  127.         DFS1(cd, root, -1);
  128.         DFS2(cd, root, -1);
  129.  
  130.         sort(all.begin(), all.end(), [&](int a, int b) {
  131.             return depth[cd][a] < depth[cd][b];
  132.         });
  133.  
  134.         for (auto x : all) {
  135.             dp[cd][x] = depth[cd][x];
  136.             if (up[x] != -1) {
  137.                 assert(depth[cd][up[x]] < depth[cd][x]);
  138.                 assert(dp[cd][up[x]] != -1);
  139.                 dp[cd][x] = dp[cd][up[x]];
  140.             }
  141.         }
  142.  
  143.         // cerr << "CENTROID: " << root + 1 << endl;
  144.  
  145.         // cerr << "SPEC: ";
  146.         // for (auto x : spec) {
  147.         //  cerr << x + 1 << " ";
  148.         // }
  149.         // cerr << endl;
  150.  
  151.         // for (auto x : all) {
  152.         //  cerr << x + 1 << ": " << endl;
  153.         //  cerr << "down: " << down[x] + 1 << endl;
  154.         //  cerr << "up: " << up[x] + 1 << endl;
  155.         //  cerr << "dp: " << dp[cd][x] << "\n";
  156.         //  // cerr << "bfsr: " << bfsr[x] << '\n';
  157.         //  // cerr << "bfsd: " << bfsd[x] << '\n';
  158.         // }
  159.         // cerr << endl;
  160.  
  161.         for (auto x : all) {
  162.             down[x] = -1;
  163.             up[x] = -1;
  164.         }
  165.     }
  166.  
  167.     void RecDecomp(int node, int size, int cp, int cd) {
  168.         // node -> centroid(node)
  169.         rec_size = size;
  170.         ComputeSizeAndCentroid(node, -1);
  171.         node = rec_centr;
  172.        
  173.         CP[node] = cp;
  174.         CD[node] = cd;
  175.  
  176.         // You can do whichever DFS you want here
  177.         all.clear();
  178.         DoUsefulDFS(node, -1, cd);
  179.         BFS(cd, node);
  180.  
  181.         for(auto vec : G[node]) {
  182.             G[vec].erase(node);
  183.             RecDecomp(vec, Size[vec], node, cd + 1);
  184.         }
  185.     }
  186.  
  187.     void Decompose(int root = 0) {
  188.         RecDecomp(root, n, -1, 0);
  189.         for(auto x : CD)
  190.             assert(x < 20);
  191.     }
  192. };
  193.  
  194. int main() {
  195.   ios_base::sync_with_stdio(false);
  196.   cin.tie(0);
  197.  
  198.   int n, k, r; cin >> n >> k >> r;
  199.   CentroidDecomposer cd(n, k);
  200.   for (int i = 1; i < n; ++i) {
  201.     int a, b; cin >> a >> b;
  202.     cd.AddEdge(a - 1, b - 1);
  203.   }
  204.   for (int i = 0; i < r; ++i) {
  205.     int x; cin >> x; cd.Tag(x - 1);
  206.   }
  207.   cd.Decompose();
  208.   int q; cin >> q;
  209.   while (q--) {
  210.     int a, b; cin >> a >> b;
  211.     cout << (cd.Solve(a - 1, b - 1) ? "YES" : "NO") << '\n';
  212.   }
  213.  
  214.   return 0;
  215. }
Add Comment
Please, Sign In to add comment