algo(dsa): quantile + greater queries on wavelet

algorithms/data-structures/wavelet-tree.cpp

@@ -23,9 +23,15 @@ class WaveletTree {
     r = new WaveletTree(pv, to, M + 1, H);
   }
 
-  // Find the k-th smallest element in positions [i,j].
-  // TO BE IMPLEMENTED
-  int quantile(int k, int i, int j) const { return 0; }
+  // Find the k-th smallest element in positions [i,j)
+  int quantile(int l, int r, int k) {
+    if (l > r) return 0;
+    if (L == H) return L;
+    int inLeft = frq[r] - frq[l - 1];
+    int lb = frq[l - 1], rb = frq[r];
+    if (k <= inLeft) return this->l->quantile(lb + 1, rb, k);
+    return this->r->quantile(l - lb, r - rb, k - inLeft);
+  }
 
   // Count occurrences of number c until position i -> [0, i].
   int rank(int c, int i) { return until(c, min(i + 1, (int)frq.size() - 1)); }
@@ -56,4 +62,17 @@ class WaveletTree {
     return this->l->range(ri, rj, a, M, L, U) +
            this->r->range(i - ri, j - rj, M + 1, b, L, U);
   }
-};
+
+  // Number of elements greater than or equal to k in [l, r];
+  // Can count distinct in a range with aux vector of next pos
+  int greater(int l, int r, int k) { return _greater(l + 1, r + 1, k); }
+
+  int _greater(int l, int r, int k) {
+    if (l > r or k > H) return 0;
+    if (L >= k) return r - l + 1;
+
+    int ri = frq[l - 1], rj = frq[r];
+    return this->l->_greater(ri + 1, rj, k) +
+           this->r->_greater(l - ri, r - rj, k);
+  }
+};