Alex_tz307

USACO 2019 December Contest, Gold Problem 2. Milk Visits - HLD

Apr 15th, 2021 (edited)
197
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #include <bits/stdc++.h>
  2.  
  3. using namespace std;
  4.  
  5. ifstream fin("milkvisits.in");
  6. ofstream fout("milkvisits.out");
  7.  
  8. const int NMAX = 1e5 + 1;
  9. int N, Q, v[NMAX], a[NMAX], x[NMAX], y[NMAX], c[NMAX], req;
  10. bool ap[NMAX];
  11. vector<int> G[NMAX];
  12. int p[NMAX], sz[NMAX], depth[NMAX], chain_top[NMAX];
  13. int cnt_labels, label[NMAX];
  14. char sol;
  15.  
  16. struct SegTree {
  17.   int Size;
  18.   vector<unordered_set<int>> tree;
  19.  
  20.   void init(int N) {
  21.     Size = 1;
  22.     while (Size < N)
  23.       Size <<= 1;
  24.     tree.resize(Size << 1 | 1);
  25.   }
  26.  
  27.   void build(int x, int lx, int rx) {
  28.     if (lx == rx) {
  29.       if (ap[a[lx]])
  30.         tree[x].emplace(a[lx]);
  31.       return;
  32.     }
  33.     int mid = (lx + rx) >> 1;
  34.     build(x << 1, lx, mid);
  35.     build(x << 1 | 1, mid + 1, rx);
  36.     for (int i = 0; i < 2; ++i)
  37.       for (int it : tree[x << 1 | i])
  38.         tree[x].emplace(it);
  39.   }
  40.  
  41.   void query(int x, int lx, int rx, int st, int dr) {
  42.     if (sol == '1')
  43.       return;
  44.     if (st <= lx && rx <= dr) {
  45.       if (tree[x].count(req))
  46.         sol = '1';
  47.       return;
  48.     }
  49.     int mid = (lx + rx) >> 1;
  50.     if (st <= mid)
  51.       query(x << 1, lx, mid, st, dr);
  52.     if (sol == '0' && mid < dr)
  53.       query(x << 1 | 1, mid + 1, rx, st, dr);
  54.   }
  55. } tree;
  56.  
  57. void read_input() {
  58.   fin >> N >> Q;
  59.   for (int i = 1; i <= N; ++i)
  60.     fin >> v[i];
  61.   for (int i = 1; i < N; ++i) {
  62.     int u, v;
  63.     fin >> u >> v;
  64.     G[u].emplace_back(v);
  65.     G[v].emplace_back(u);
  66.   }
  67.   for (int q = 0; q < Q; ++q) {
  68.     fin >> x[q] >> y[q] >> c[q];
  69.     ap[c[q]] = true;
  70.   }
  71. }
  72.  
  73. void dfs_init(int u, int parent) {
  74.   p[u] = parent;
  75.   depth[u] = depth[parent] + 1;
  76.   sz[u] = 1;
  77.   chain_top[u] = u;
  78.   for (int v : G[u])
  79.     if (v != parent) {
  80.       dfs_init(v, u);
  81.       sz[u] += sz[v];
  82.     }
  83. }
  84.  
  85. void dfs_heavy(int u) {
  86.   label[u] = ++cnt_labels;
  87.   int max_size = -1, heavy_son = -1;
  88.   for (int v : G[u])
  89.     if (v != p[u] && sz[v] > max_size) {
  90.       max_size = sz[v];
  91.       heavy_son = v;
  92.     }
  93.   if (heavy_son == -1)
  94.     return;
  95.   chain_top[heavy_son] = chain_top[u];
  96.   dfs_heavy(heavy_son);
  97.   for (int v : G[u])
  98.     if (v != p[u] && v != heavy_son)
  99.       dfs_heavy(v);
  100. }
  101.  
  102. void lca_query(int x, int y) {
  103.   if (sol == '1')
  104.     return;
  105.   if (chain_top[x] == chain_top[y]) {
  106.     if (label[x] > label[y])
  107.       swap(x, y);
  108.     tree.query(1, 1, N, label[x], label[y]);
  109.     return;
  110.   }
  111.   if (depth[p[chain_top[x]]] > depth[p[chain_top[y]]]) {
  112.     tree.query(1, 1, N, label[chain_top[x]], label[x]);
  113.     if (sol == '0')
  114.       lca_query(p[chain_top[x]], y);
  115.     return;
  116.   }
  117.   tree.query(1, 1, N, label[chain_top[y]], label[y]);
  118.   if (sol == '0')
  119.     lca_query(x, p[chain_top[y]]);
  120. }
  121.  
  122. void solve() {
  123.   dfs_init(1, 0);
  124.   dfs_heavy(1);
  125.   tree.init(N);
  126.   for (int i = 1; i <= N; ++i)
  127.     a[label[i]] = v[i];
  128.   tree.build(1, 1, N);
  129.   for (int q = 0; q < Q; ++q) {
  130.     req = c[q];
  131.     sol = '0';
  132.     lca_query(x[q], y[q]);
  133.     fout << sol;
  134.   }
  135. }
  136.  
  137. void close_files() {
  138.   fin.close();
  139.   fout.close();
  140. }
  141.  
  142. int main() {
  143.   read_input();
  144.   solve();
  145.   close_files();
  146.   return 0;
  147. }
  148.  
RAW Paste Data