code/dataStructures/WLTsmall.cc

View this file on GitHub
Last update: 2020-12-19 23:25:36+01:00

Depends on

code/template.cc

Code

#include "../template.cc"
// only suitable for relatively small alphabets ~10^6
template<typename T = int>
struct WLT {
  using vt = typename vector<T>::iterator;
  T H;
  vvi M;
  static inline T mid(T a, T b) { return (a + b + 1) / 2; }
  static inline int left(int i) { return i << 1; }
  static inline int right(int i) { return i << 1 | 1; }
  // [0, hi)
  WLT(vector<T> v, T hi) : H{hi}, M(H) {
    build(ALL(v), 1, 0, H);
  }
  void build(vt b, vt e, int i, T lo, T hi) {
    if (lo + 1 == hi) return;
    int m = mid(lo, hi);
    M[i].reserve(distance(b, e) + 1); M[i].pb(0);
    auto less = [=](T k) { return k < m; };
    for (auto it = b; it != e; ++it) M[i].pb(M[i].back() + less(*it));
    auto p = stable_partition(b, e, less);
    if (p != b) build(b, p, left(i), lo, m); 
    if (p != e) build(p, e, right(i), m, hi);
  }
  // k-th smallest in [l, r)
  T kthsmallest(int l, int r, int k) {
    T lo = 0, hi = H; int i = 1;
    while (lo + 1 != hi) {
      if (k < M[i][r] - M[i][l]) {
        l = M[i][l]; r = M[i][r];
        i = left(i); hi = mid(lo, hi);
      } else {
        k -= M[i][r] - M[i][l];
        l -= M[i][l]; r -= M[i][r];
        i = right(i); lo = mid(lo, hi);
      }
    }
    return lo;
  }
};

#line 1 "code/template.cc"
// this line is here for a reason
#include <bits/stdc++.h>
using namespace std;
typedef long long ll;
typedef pair<int, int> ii;
typedef vector<int> vi;
typedef vector<ii> vii;
typedef vector<vi> vvi;
typedef vector<vii> vvii;
#define fi first
#define se second
#define eb emplace_back
#define pb push_back
#define mp make_pair
#define mt make_tuple
#define endl '\n'
#define ALL(x) (x).begin(), (x).end()
#define RALL(x) (x).rbegin(), (x).rend()
#define SZ(x) (int)(x).size()
#define FOR(a, b, c) for (auto a = (b); (a) < (c); ++(a))
#define F0R(a, b) FOR (a, 0, (b))
template <typename T>
bool ckmin(T& a, const T& b) { return a > b ? a = b, true : false; }
template <typename T>
bool ckmax(T& a, const T& b) { return a < b ? a = b, true : false; }
#ifndef DEBUG
#define DEBUG 0
#endif
#define dout if (DEBUG) cerr
#define dvar(...) " [" << #__VA_ARGS__ ": " << (__VA_ARGS__) << "] "
#line 2 "code/dataStructures/WLTsmall.cc"
// only suitable for relatively small alphabets ~10^6
template<typename T = int>
struct WLT {
  using vt = typename vector<T>::iterator;
  T H;
  vvi M;
  static inline T mid(T a, T b) { return (a + b + 1) / 2; }
  static inline int left(int i) { return i << 1; }
  static inline int right(int i) { return i << 1 | 1; }
  // [0, hi)
  WLT(vector<T> v, T hi) : H{hi}, M(H) {
    build(ALL(v), 1, 0, H);
  }
  void build(vt b, vt e, int i, T lo, T hi) {
    if (lo + 1 == hi) return;
    int m = mid(lo, hi);
    M[i].reserve(distance(b, e) + 1); M[i].pb(0);
    auto less = [=](T k) { return k < m; };
    for (auto it = b; it != e; ++it) M[i].pb(M[i].back() + less(*it));
    auto p = stable_partition(b, e, less);
    if (p != b) build(b, p, left(i), lo, m); 
    if (p != e) build(p, e, right(i), m, hi);
  }
  // k-th smallest in [l, r)
  T kthsmallest(int l, int r, int k) {
    T lo = 0, hi = H; int i = 1;
    while (lo + 1 != hi) {
      if (k < M[i][r] - M[i][l]) {
        l = M[i][l]; r = M[i][r];
        i = left(i); hi = mid(lo, hi);
      } else {
        k -= M[i][r] - M[i][l];
        l -= M[i][l]; r -= M[i][r];
        i = right(i); lo = mid(lo, hi);
      }
    }
    return lo;
  }
};