Alex_tz307

Wavelet Tree

Apr 23rd, 2021 (edited)
253
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 1.50 KB | None | 0 0
  1. #include <bits/stdc++.h>
  2.  
  3. using namespace std;
  4.  
  5. void fastIO() {
  6.   ios_base::sync_with_stdio(false);
  7.   cin.tie(nullptr);
  8.   cout.tie(nullptr);
  9. }
  10.  
  11. typedef vector<int>::iterator itr;
  12.  
  13. struct wavelet {
  14.   wavelet *l, *r;
  15.   vector<int> f;
  16.   int lo, hi;
  17.  
  18.   void build(itr st, itr dr, int x, int y) {
  19.     lo = x, hi = y;
  20.     if (x == y)
  21.       return;
  22.     f.reserve(dr - st + 1);
  23.     int mid = (x + y) >> 1;
  24.     auto check = [&](int v) {
  25.       return v <= mid;
  26.     };
  27.     f.push_back(0);
  28.     for (itr it = st; it != dr; ++it)
  29.       f.push_back(f.back() + check(*it));
  30.     itr pivot = stable_partition(st, dr, check);
  31.     if (st < pivot) {
  32.       l = new wavelet;
  33.       l->build(st, pivot, x, mid);
  34.     }
  35.     if (pivot < dr) {
  36.       r = new wavelet;
  37.       r->build(pivot, dr, mid + 1, y);
  38.     }
  39.   }
  40.  
  41.   int kth(int st, int dr, int k) {
  42.     if (lo == hi)
  43.       return lo;
  44.     int lb = f[st - 1], rb = f[dr];
  45.     int in_left = rb - lb;
  46.     if (k <= in_left)
  47.       return l->kth(lb + 1, rb, k);
  48.     return r->kth(st - lb, dr - rb, k - in_left);
  49.   }
  50. };
  51.  
  52. void solve() {
  53.   int N, Q;
  54.   cin >> N >> Q;
  55.   vector<int> a(N + 1);
  56.   for (int i = 1; i <= N; ++i)
  57.     cin >> a[i];
  58.   wavelet wt;
  59.   wt.build(a.begin() + 1, a.end(), *min_element(a.begin() + 1, a.end()), *max_element(a.begin() + 1, a.end()));
  60.   for (int i = 0; i < Q; ++i) {
  61.     int l, r, k;
  62.     cin >> l >> r >> k;
  63.     ++l, ++k;
  64.     cout << wt.kth(l, r, k) << '\n';
  65.   }
  66. }
  67.  
  68. int main() {
  69.   fastIO();
  70.   solve();
  71.   return 0;
  72. }
  73.  
Add Comment
Please, Sign In to add comment