AhmedAshraff

Untitled

Nov 26th, 2024 (edited)
40
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 1.95 KB | None | 0 0
  1. class WaveletTree {// 0-based , values of arr better to be compressed
  2. public:
  3.     int low, high;
  4.     WaveletTree *left, *right;
  5.     vector<int> freq;
  6.  
  7.     WaveletTree(vector<int>& arr, int l, int r, int lo, int hi) : low(lo), high(hi), left(nullptr), right(nullptr) {
  8.         if (lo == hi || l > r) return;
  9.  
  10.         int mid = (lo + hi) / 2;
  11.         freq.reserve(r - l + 2);
  12.         freq.push_back(0);
  13.  
  14.         for (int i = l; i <= r; i++) {
  15.             freq.push_back(freq.back() + (arr[i] <= mid));
  16.         }
  17.  
  18.         vector<int> leftArr, rightArr;
  19.         for (int i = l; i <= r; i++) {
  20.             if (arr[i] <= mid)
  21.                 leftArr.push_back(arr[i]);
  22.             else
  23.                 rightArr.push_back(arr[i]);
  24.         }
  25.  
  26.         if (!leftArr.empty()) {
  27.             left = new WaveletTree(leftArr, 0, sz(leftArr) - 1, lo, mid);
  28.         }
  29.         if (!rightArr.empty()) {
  30.             right = new WaveletTree(rightArr, 0, sz(rightArr) - 1, mid + 1, hi);
  31.         }
  32.     }
  33.  
  34.         int kth(int l, int r, int k) {
  35.         if (low == high) return low;
  36.         int inLeft = freq[r + 1] - freq[l];
  37.         if (k <= inLeft) {
  38.             return left->kth(freq[l], freq[r + 1] - 1, k);
  39.         } else {
  40.             return right->kth(l - freq[l], r - freq[r + 1], k - inLeft);
  41.         }
  42.     }
  43.  
  44.     int lte(int l, int r, int x) {
  45.         if (x < low) return 0;
  46.         if (x >= high) return r - l + 1;
  47.         return left->lte(freq[l], freq[r + 1] - 1, x) + right->lte(l - freq[l], r - freq[r + 1], x);
  48.     }
  49.  
  50.     int count(int l, int r, int x) {
  51.         if (low == high) return (low == x) ? r - l + 1 : 0;
  52.         int mid = (low + high) / 2;
  53.         if (x <= mid) {
  54.             return left ? left->count(freq[l], freq[r + 1] - 1, x) : 0;
  55.         } else {
  56.             return right ? right->count(l - freq[l], r - freq[r + 1], x) : 0;
  57.         }
  58.     }
  59.  
  60.     ~WaveletTree() {
  61.         delete left;
  62.         delete right;
  63.     }
  64. };
Advertisement
Add Comment
Please, Sign In to add comment