fork download 
  1. #include <bits/stdc++.h>
  2. using namespace std;
  3. #define fo(i,n) for(i=0;i<n;i++)
  4. #define ll long long
  5.  
  6. #define pb push_back
  7. #define mp make_pair
  8.  
  9.  
  10. typedef pair<int, int>	pii;
  11. typedef pair<ll, ll>	pl;
  12. typedef vector<int>		vi;
  13.  
  14. const int N = 3e5, M = N;
  15. //=======================
  16. const int MAX = 1e6;
  17.  
  18. int a[N];
  19. struct wavelet_tree
  20. {
  21. #define vi vector<int>
  22. #define pb push_back
  23.     int lo, hi;
  24.     wavelet_tree *l=0, *r=0;
  25.     vi b;
  26.     vi c; // c holds the prefix sum of elements
  27.  
  28.     //nos are in range [x,y]
  29.     //array indices are [from, to)
  30.     wavelet_tree(int *from, int *to, int x, int y)
  31.     {
  32.         lo = x, hi = y;
  33.         if( from >= to)
  34.             return;
  35.         if( hi == lo )
  36.         {
  37.             b.reserve(to-from+1);
  38.             b.pb(0);
  39.             c.reserve(to-from+1);
  40.             c.pb(0);
  41.             for(auto it = from; it != to; it++)
  42.             {
  43.                 b.pb(b.back() + 1);
  44.                 c.pb(c.back()+*it);
  45.             }
  46.             return ;
  47.         }
  48.         int mid = (lo+hi)/2;
  49.         auto f = [mid](int x)
  50.         {
  51.             return x <= mid;
  52.         };
  53.         b.reserve(to-from+1);
  54.         b.pb(0);
  55.         c.reserve(to-from+1);
  56.         c.pb(0);
  57.         for(auto it = from; it != to; it++)
  58.         {
  59.             b.pb(b.back() + f(*it));
  60.             c.pb(c.back() + *it);
  61.         }
  62.         //see how lambda function is used here
  63.         auto pivot = stable_partition(from, to, f);
  64.         l = new wavelet_tree(from, pivot, lo, mid);
  65.         r = new wavelet_tree(pivot, to, mid+1, hi);
  66.     }
  67.  
  68.     // swap a[i] with a[i+1]  , if a[i]!=a[i+1] call swapadjacent(i)
  69.     void swapadjacent(int i)
  70.     {
  71.         if(lo == hi)
  72.             return ;
  73.         b[i]= b[i-1] + b[i+1] - b[i];
  74.         c[i] = c[i-1] + c[i+1] - c[i];
  75.         if( b[i+1]-b[i] == b[i] - b[i-1])
  76.         {
  77.             if(b[i]-b[i-1])
  78.                 return this->l->swapadjacent(b[i]);
  79.             else
  80.                 return this->r->swapadjacent(i-b[i]);
  81.         }
  82.         else
  83.             return ;
  84.     }
  85.  
  86.     //kth smallest element in [l, r]
  87.     int kth(int l, int r, int k)
  88.     {
  89.         if(l > r)
  90.             return 0;
  91.         if(lo == hi)
  92.             return lo;
  93.         int inLeft = b[r] - b[l-1];
  94.         int lb = b[l-1]; //amt of nos in first (l-1) nos that go in left
  95.         int rb = b[r]; //amt of nos in first (r) nos that go in left
  96.         if(k <= inLeft)
  97.             return this->l->kth(lb+1, rb, k);
  98.         return this->r->kth(l-lb, r-rb, k-inLeft);
  99.     }
  100.  
  101.     //count of nos in [l, r] Less than or equal to k
  102.     int LTE(int l, int r, int k)
  103.     {
  104.         if(l > r or k < lo)
  105.             return 0;
  106.         if(hi <= k)
  107.             return r - l + 1;
  108.         int lb = b[l-1], rb = b[r];
  109.         return this->l->LTE(lb+1, rb, k) + this->r->LTE(l-lb, r-rb, k);
  110.     }
  111.  
  112.     //count of nos in [l, r] equal to k
  113.     int count(int l, int r, int k)
  114.     {
  115.         if(l > r or k < lo or k > hi)
  116.             return 0;
  117.         if(lo == hi)
  118.             return r - l + 1;
  119.         int lb = b[l-1], rb = b[r], mid = (lo+hi)/2;
  120.         if(k <= mid)
  121.             return this->l->count(lb+1, rb, k);
  122.         return this->r->count(l-lb, r-rb, k);
  123.     }
  124.  
  125.     //sum of nos in [l ,r] less than or equal to k
  126.     int sumk(int l, int r, int k)
  127.     {
  128.         if(l > r or k < lo)
  129.             return 0;
  130.         if(hi <= k)
  131.             return c[r] - c[l-1];
  132.         int lb = b[l-1], rb = b[r];
  133.         return this->l->sumk(lb+1, rb, k) + this->r->sumk(l-lb, r-rb, k);
  134.     }
  135.  
  136.     ~wavelet_tree()
  137.     {
  138.         if(l)
  139.         delete l;
  140.         if(r)
  141.         delete r;
  142.     }
  143. };
  144. int main()
  145. {
  146.     ios_base::sync_with_stdio(false);
  147.     cin.tie(NULL);
  148.     srand(time(NULL));
  149.     int i,n,k,j,q,l,r;
  150.     cin >> n;
  151.     fo(i, n) cin >> a[i+1];
  152.     wavelet_tree T(a+1, a+n+1, 1, MAX);
  153.     cin >> q;
  154.     while(q--)
  155.     {
  156.         int x;
  157.         cin >> x;
  158.         cin >> l >> r >> k;
  159.         if(x == 0)
  160.         {
  161.             //kth smallest
  162.             cout << "Kth smallest: ";
  163.             cout << T.kth(l, r, k) << endl;
  164.         }
  165.         if(x == 1)
  166.         {
  167.             //less than or equal to K
  168.             cout << "LTE: ";
  169.             cout << T.LTE(l, r, k) << endl;
  170.         }
  171.         if(x == 2)
  172.         {
  173.             //count occurence of K in [l, r]
  174.             cout << "Occurence of K: ";
  175.             cout << T.count(l, r, k) << endl;
  176.         }
  177.         if(x == 3)
  178.         {
  179.             //sum of elements less than or equal to K in [l, r]
  180.             cout << "Sum: ";
  181.             cout << T.sumk(l, r, k) << endl;
  182.         }
  183.  
  184.     }
  185.     return 0;
  186. }
  187.  
  188.  
Time limit exceeded #stdin #stdout 5s 4404KB
comments (?)
stdin 
Standard input is empty
stdout 
Standard output is empty
https://ideone.com/ARDx02
language:
C++ 4.3.2 (gcc-4.3.2)
created:
5 years ago
visibility:
secret

Share or Embed source code

Discover > Sphere Engine API

The brand new service which powers Ideone!

Discover > IDE Widget

Widget for compiling and running the source code in a web browser!