Advertisement
dipBRUR

19

Nov 19th, 2017
100
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
text 1.71 KB | None | 0 0
  1. // Algo : Wavelet Tree
  2.  
  3. struct wavelet_tree{
  4. #define vi vector<int>
  5. #define pb push_back
  6.  
  7. int lo, hi;
  8. wavelet_tree *l, *r;
  9. vi b;
  10.  
  11. //nos are in range [x,y]
  12. //array indices are [from, to)
  13. wavelet_tree(int *from, int *to, int x, int y){
  14. lo = x, hi = y;
  15. if(lo == hi or from >= to) return;
  16. int mid = (lo+hi)/2;
  17. auto f = [mid](int x){
  18. return x <= mid;
  19. };
  20. b.reserve(to-from+1);
  21. b.pb(0);
  22. for(auto it = from; it != to; it++)
  23. b.pb(b.back() + f(*it));
  24. //see how lambda function is used here
  25. auto pivot = stable_partition(from, to, f);
  26. l = new wavelet_tree(from, pivot, lo, mid);
  27. r = new wavelet_tree(pivot, to, mid+1, hi);
  28. }
  29.  
  30. //kth smallest element in [l, r]
  31. int kth(int l, int r, int k){
  32. if(l > r) return 0;
  33. if(lo == hi) return lo;
  34. int inLeft = b[r] - b[l-1];
  35. int lb = b[l-1]; //amt of nos in first (l-1) nos that go in left
  36. int rb = b[r]; //amt of nos in first (r) nos that go in left
  37. if(k <= inLeft) return this->l->kth(lb+1, rb , k);
  38. return this->r->kth(l-lb, r-rb, k-inLeft);
  39. }
  40.  
  41. //count of nos in [l, r] Less than or equal to k
  42. int LTE(int l, int r, int k) {
  43. if(l > r or k < lo) return 0;
  44. if(hi <= k) return r - l + 1;
  45. int lb = b[l-1], rb = b[r];
  46. return this->l->LTE(lb+1, rb, k) + this->r->LTE(l-lb, r-rb, k);
  47. }
  48.  
  49. //count of nos in [l, r] equal to k
  50. int count(int l, int r, int k) {
  51. if(l > r or k < lo or k > hi) return 0;
  52. if(lo == hi) return r - l + 1;
  53. int lb = b[l-1], rb = b[r], mid = (lo+hi)/2;
  54. if(k <= mid) return this->l->count(lb+1, rb, k);
  55. return this->r->count(l-lb, r-rb, k);
  56. }
  57. ~wavelet_tree(){
  58. delete l;
  59. delete r;
  60. }
  61. };
  62.  
  63. // Call :
  64. wavelet_tree *T;
  65. T = new wavelet_tree(arr+1, arr+N+1, 1, MAXN);
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement