Skip to the content.

:heavy_check_mark: data-structure/wavelet-matrix-with-segment-tree.hpp

Depends on

Verified with

Code

#pragma once

#include "data-structure/wavelet-matrix.hpp"
#include "segment-tree/segment-tree.hpp"

// M: commutative monoid
template <class T, class M, int B = 30>
struct WaveletMatrixWithSegmentTree : public WaveletMatrix<T, B> {
  using Base = WaveletMatrix<T, B>;
  using S = typename M::value_type;
  using u32 = uint32_t;
  using i64 = int64_t;
  using u64 = uint64_t;

  using Base::a;
  using Base::bv;
  using Base::n;
  vector<S> w;
  vector<SegmentTree<M>> seg;

  WaveletMatrixWithSegmentTree(u32 _n) : Base(_n), w(_n) {}
  WaveletMatrixWithSegmentTree(const vector<T>& _a, const vector<S>& _w) : Base(_a), w(_w) { build(); }

  void set(u32 i, const T& x, const S& v) {
    assert(x >= 0);
    a[i] = x;
    w[i] = v;
  }

  void build() {
    bv.assign(B, n);
    seg.resize(B + 1);
    seg[B] = SegmentTree<M>(w);
    vector<T> cur = a, nxt(n);
    vector<S> wcur = w, wnxt(n);
    for (int h = B - 1; h >= 0; --h) {
      for (int i = 0; i < n; ++i)
        if ((cur[i] >> h) & 1) bv[h].set(i);
      bv[h].build();
      array<decltype(begin(nxt)), 2> it{begin(nxt), begin(nxt) + bv[h].zeros};
      array<decltype(begin(wnxt)), 2> wit{begin(wnxt), begin(wnxt) + bv[h].zeros};
      for (int i = 0; i < n; ++i) {
        int x = bv[h].get(i);
        *it[x]++ = cur[i];
        *wit[x]++ = wcur[i];
      }
      seg[h] = SegmentTree<M>(wnxt);
      swap(cur, nxt);
      swap(wcur, wnxt);
    }
  }

  void update(u32 k, const S& v) {
    seg[B].set(k, v);
    for (int h = B - 1; h >= 0; --h) {
      u32 f = bv[h].get(k);
      k = f ? bv[h].rank1(k) + bv[h].zeros : bv[h].rank0(k);
      seg[h].set(k, v);
    }
  }
  void apply(u32 k, const S& v) {
    seg[B].apply(k, v);
    for (int h = B - 1; h >= 0; --h) {
      u32 f = bv[h].get(k);
      k = f ? bv[h].rank1(k) + bv[h].zeros : bv[h].rank0(k);
      seg[h].apply(k, v);
    }
  }

  // count i s.t. (l <= i < r) && (lower <= v[i] ^ value_xor < upper)
  S range_sum(int l, int r, T lower, T upper, T value_xor = 0) {
    if (lower >= upper) return M::e();
    return range_sum_(B - 1, l, r, T(0), T(1) << B, lower, upper, value_xor);
  }

 private:
  S range_sum_(int h, int l, int r, T vl, T vr, T lower, T upper, T value_xor) {
    if (r <= l) return M::e();
    if (vr <= lower || upper <= vl) return M::e();
    if (lower <= vl && vr <= upper) return seg[h + 1].prod(l, r);

    u32 l0 = bv[h].rank0(l), r0 = bv[h].rank0(r);
    u32 zeros = bv[h].zeros;
    u32 l1 = l + zeros - l0, r1 = r + zeros - r0;
    if ((value_xor >> h) & 1) {
      swap(l0, l1);
      swap(r0, r1);
    }

    T vm = (vl + vr) >> 1;
    return M::op(range_sum_(h - 1, l0, r0, vl, vm, lower, upper, value_xor),
                 range_sum_(h - 1, l1, r1, vm, vr, lower, upper, value_xor));
  }
};
#line 2 "data-structure/wavelet-matrix-with-segment-tree.hpp"

#line 2 "data-structure/wavelet-matrix.hpp"

#line 2 "data-structure/bit-vector.hpp"

struct BitVector {
  using i32 = int32_t;
  using u32 = uint32_t;
  using u64 = uint64_t;

  static constexpr u32 W = 64;
  inline u32 get(u32 i) const { return u32(block[i / W] >> (i % W)) & 1u; }
  inline void set(u32 i) { block[i / W] |= 1ull << (i % W); }

  vector<u64> block;
  vector<i32> count;
  u32 n, zeros;
  BitVector() {}
  BitVector(int _n) : n(_n) {
    block.resize(n / W + 1, 0);
    count.resize(block.size(), 0);
  }
  void build() {
    for (u32 i = 1; i < block.size(); i++)
      count[i] = count[i - 1] + __builtin_popcountll(block[i - 1]);
    zeros = rank0(n);
  }
  inline u32 rank0(u32 i) const { return i - rank1(i); }
  inline u32 rank1(u32 i) const { return count[i / W] + __builtin_popcountll(block[i / W] & ((1ull << i % W) - 1)); }
};
#line 4 "data-structure/wavelet-matrix.hpp"

template <class T, int B = 30>
struct WaveletMatrix {
  using u32 = uint32_t;
  using i64 = int64_t;
  using u64 = uint64_t;

  int n;
  vector<T> a;
  vector<BitVector> bv;

  WaveletMatrix(u32 _n) : n(max<u32>(_n, 1)), a(n) {}
  WaveletMatrix(const vector<T>& _a) : n(_a.size()), a(_a) { build(); }

  void set(u32 i, const T& x) {
    assert(x >= 0);
    a[i] = x;
  }

  void build() {
    bv.assign(B, n);
    vector<T> cur = a, nxt(n);
    for (int h = B - 1; h >= 0; --h) {
      for (int i = 0; i < n; ++i)
        if ((cur[i] >> h) & 1) bv[h].set(i);
      bv[h].build();
      array<decltype(begin(nxt)), 2> it{begin(nxt), begin(nxt) + bv[h].zeros};
      for (int i = 0; i < n; ++i) *it[bv[h].get(i)]++ = cur[i];
      swap(cur, nxt);
    }
  }

  inline pair<u32, u32> succ0(int l, int r, int h) const {
    return make_pair(bv[h].rank0(l), bv[h].rank0(r));
  }
  inline pair<u32, u32> succ1(int l, int r, int h) const {
    u32 l0 = bv[h].rank0(l);
    u32 r0 = bv[h].rank0(r);
    u32 zeros = bv[h].zeros;
    return make_pair(l + zeros - l0, r + zeros - r0);
  }

  // return a[k]
  T access(u32 k) const {
    T ret = 0;
    for (int h = B - 1; h >= 0; --h) {
      u32 f = bv[h].get(k);
      ret |= f ? T(1) << h : 0;
      k = f ? bv[h].rank1(k) + bv[h].zeros : bv[h].rank0(k);
    }
    return ret;
  }

  // k-th (0-indexed) smallest number in { a[i] ^ value_xor : i in [l, r) }
  T kth_smallest(u32 l, u32 r, u32 k, T value_xor = 0) const {
    T res = value_xor;
    for (int h = B - 1; h >= 0; --h) {
      u32 l0 = bv[h].rank0(l), r0 = bv[h].rank0(r);
      u32 c0 = r0 - l0;
      if ((k < c0) ^ ((value_xor >> h) & 1))
        l = l0, r = r0;
      else {
        k -= c0;
        res ^= (T)1 << h;
        l += bv[h].zeros - l0;
        r += bv[h].zeros - r0;
      }
    }
    return res;
  }
  // k-th (0-indexed) largest number in { a[i] ^ value_xor : i in [l, r) }
  T kth_largest(u32 l, u32 r, u32 k, T value_xor = 0) {
    return kth_smallest(l, r, r - l - k - 1);
  }

  // count i s.t. (l <= i < r) && (v[i] ^ value_xor < upper)
  int range_freq(int l, int r, T upper, T value_xor = 0) {
    if (upper >= (T(1) << B)) return r - l;
    int ret = 0;
    for (int h = B - 1; h >= 0; --h) {
      bool f = (upper >> h) & 1;
      u32 l0 = bv[h].rank0(l), r0 = bv[h].rank0(r);
      u32 zeros = bv[h].zeros;
      u32 l1 = l + zeros - l0, r1 = r + zeros - r0;
      if ((value_xor >> h) & 1) {
        swap(l0, l1);
        swap(r0, r1);
      }
      if (f) {
        ret += r0 - l0;
        l += zeros - l0;
        r += zeros - r0;
      } else {
        l = l0;
        r = r0;
      }
    }
    return ret;
  }
  int range_freq(int l, int r, T lower, T upper, T value_xor) {
    return range_freq(l, r, upper, value_xor) - range_freq(l, r, lower, value_xor);
  }

  // max v[i] s.t. (l <= i < r) && (v[i] ^ value_xor < upper)
  T prev_value(int l, int r, T upper, T value_xor = 0) {
    int cnt = range_freq(l, r, upper, value_xor);
    return cnt == 0 ? T(-1) : kth_smallest(l, r, cnt - 1, value_xor);
  }

  // min v[i] s.t. (l <= i < r) && (lower ^ value_xor <= v[i])
  T next_value(int l, int r, T lower, T value_xor = 0) {
    int cnt = range_freq(l, r, lower, value_xor);
    return cnt == r - l ? T(-1) : kth_smallest(l, r, cnt, value_xor);
  }
};

/**
 * @brief Wavelet Matrix
 * @docs docs/data-structure/wavelet-matrix.md
 */
#line 2 "algebraic-structure/util.hpp"
#ifdef __cpp_concepts
#define REQUIRES(...) requires __VA_ARGS__
#else
#define REQUIRES(...)
#endif
#line 3 "algebraic-structure/magma.hpp"

#ifdef __cpp_concepts
template <class M>
concept Magma = requires(typename M::value_type x, typename M::value_type y) {
  typename M::value_type;
  { M::op(x, y) } -> same_as<typename M::value_type>;
};
#endif

template <class T>
struct AddMagma {
  using value_type = T;
  static T op(T x, T y) { return x + y; }
};
template <class T>
struct MulMagma {
  using value_type = T;
  static T op(T x, T y) { return x * y; }
};
template <class T, T id>
struct MaxMagma {
  using value_type = T;
  static T op(T x, T y) { return x > y ? x : y; }
};
template <class T, T id>
struct MinMagma {
  using value_type = T;
  static T op(T x, T y) { return x < y ? x : y; }
};
#line 3 "algebraic-structure/monoid.hpp"

#ifdef __cpp_concepts
template <class M>
concept Monoid = Magma<M> && requires {
  { M::e() } -> same_as<typename M::value_type>;
};
#endif

template <class T>
struct AddMonoid {
  using value_type = T;
  static T op(T x, T y) { return x + y; }
  static T e() { return T(0); }
};
template <class T>
struct MulMonoid {
  using value_type = T;
  static T op(T x, T y) { return x * y; }
  static T e() { return T(1); }
};
template <class T, T id>
struct MaxMonoid {
  using value_type = T;
  static T op(T x, T y) { return x > y ? x : y; }
  static T e() { return id; }
};
template <class T, T id>
struct MinMonoid {
  using value_type = T;
  static T op(T x, T y) { return x < y ? x : y; }
  static T e() { return id; }
};
#line 3 "segment-tree/segment-tree.hpp"

template <class M>
REQUIRES(Monoid<M>)
struct SegmentTree {
  using T = typename M::value_type;

 private:
  int _n, size, log;
  vector<T> d;
  void update(int p) { d[p] = M::op(d[2 * p], d[2 * p + 1]); }

 public:
  SegmentTree() : SegmentTree(0) {}
  explicit SegmentTree(int sz) : SegmentTree(vector<T>(sz, M::e())) {}
  explicit SegmentTree(const vector<T>& v) : _n(v.size()) {
    size = 1, log = 0;
    while (size < _n) size <<= 1, log++;
    d.assign(2 * size, M::e());
    for (int i = 0; i < _n; i++) d[size + i] = v[i];
    for (int i = size - 1; i > 0; i--) update(i);
  }
  void clear() { fill(d.begin(), d.end(), M::e()); }

  void set_without_update(int p, T v) { d[p + size] = v; }
  void all_update() {
    for (int i = size - 1; i > 0; i--) update(i);
  }
  T get(int p) {
    assert(0 <= p && p <= _n);
    return d[p + size];
  }
  void set(int p, T v) {
    assert(0 <= p && p <= _n);
    p += size;
    d[p] = v;
    for (int i = 1; i <= log; i++) update(p >> i);
  }
  void apply(int p, T v) {
    assert(0 <= p && p <= _n);
    p += size;
    d[p] = M::op(d[p], v);
    for (int i = 1; i <= log; i++) update(p >> i);
  }
  T all_prod() { return d[1]; }
  T prod(int l, int r) {
    if (l >= r) return M::e();
    assert(0 <= l && l <= r && r <= _n);
    T sl = M::e(), sr = M::e();
    l += size, r += size;
    while (l < r) {
      if ((l & 1) != 0) sl = M::op(sl, d[l++]);
      if ((r & 1) != 0) sr = M::op(d[--r], sr);
      l >>= 1, r >>= 1;
    }
    return M::op(sl, sr);
  }

  template <bool (*f)(T)>
  int max_right(int l) const {
    return max_right(l, [](T x) { return f(x); });
  }
  template <class F>
  int max_right(int l, F f) const {
    assert(0 <= l && l <= size);
    assert(f(M::e()));
    if (l == _n) return _n;
    l += size;
    T s = M::e();
    do {
      while (l % 2 == 0) l >>= 1;
      if (!f(M::op(s, d[l]))) {
        while (l < size) {
          l <<= 1;
          if (f(M::op(s, d[l]))) s = M::op(s, d[l++]);
        }
        return l - size;
      }
      s = M::op(s, d[l++]);
    } while ((l & -l) != l);
    return _n;
  }

  template <bool (*f)(T)>
  int min_left(int r) const {
    return min_left(r, [](T x) { return f(x); });
  }
  template <class F>
  int min_left(int r, F f) const {
    assert(0 <= r && r <= _n);
    assert(f(M::e()));
    if (r == 0) return 0;
    r += size;
    T s = M::e();
    do {
      r--;
      while (r > 1 && (r % 2)) r >>= 1;
      if (!f(M::op(d[r], s))) {
        while (r < size) {
          r <<= 1, r++;
          if (f(M::op(d[r], s))) s = M::op(d[r--], s);
        }
        return r + 1 - size;
      }
      s = M::op(d[r], s);
    } while ((r & -r) != r);
    return 0;
  }
};

/**
 * @brief Segment Tree
 * @docs docs/segment-tree/segment-tree.md
 */
#line 5 "data-structure/wavelet-matrix-with-segment-tree.hpp"

// M: commutative monoid
template <class T, class M, int B = 30>
struct WaveletMatrixWithSegmentTree : public WaveletMatrix<T, B> {
  using Base = WaveletMatrix<T, B>;
  using S = typename M::value_type;
  using u32 = uint32_t;
  using i64 = int64_t;
  using u64 = uint64_t;

  using Base::a;
  using Base::bv;
  using Base::n;
  vector<S> w;
  vector<SegmentTree<M>> seg;

  WaveletMatrixWithSegmentTree(u32 _n) : Base(_n), w(_n) {}
  WaveletMatrixWithSegmentTree(const vector<T>& _a, const vector<S>& _w) : Base(_a), w(_w) { build(); }

  void set(u32 i, const T& x, const S& v) {
    assert(x >= 0);
    a[i] = x;
    w[i] = v;
  }

  void build() {
    bv.assign(B, n);
    seg.resize(B + 1);
    seg[B] = SegmentTree<M>(w);
    vector<T> cur = a, nxt(n);
    vector<S> wcur = w, wnxt(n);
    for (int h = B - 1; h >= 0; --h) {
      for (int i = 0; i < n; ++i)
        if ((cur[i] >> h) & 1) bv[h].set(i);
      bv[h].build();
      array<decltype(begin(nxt)), 2> it{begin(nxt), begin(nxt) + bv[h].zeros};
      array<decltype(begin(wnxt)), 2> wit{begin(wnxt), begin(wnxt) + bv[h].zeros};
      for (int i = 0; i < n; ++i) {
        int x = bv[h].get(i);
        *it[x]++ = cur[i];
        *wit[x]++ = wcur[i];
      }
      seg[h] = SegmentTree<M>(wnxt);
      swap(cur, nxt);
      swap(wcur, wnxt);
    }
  }

  void update(u32 k, const S& v) {
    seg[B].set(k, v);
    for (int h = B - 1; h >= 0; --h) {
      u32 f = bv[h].get(k);
      k = f ? bv[h].rank1(k) + bv[h].zeros : bv[h].rank0(k);
      seg[h].set(k, v);
    }
  }
  void apply(u32 k, const S& v) {
    seg[B].apply(k, v);
    for (int h = B - 1; h >= 0; --h) {
      u32 f = bv[h].get(k);
      k = f ? bv[h].rank1(k) + bv[h].zeros : bv[h].rank0(k);
      seg[h].apply(k, v);
    }
  }

  // count i s.t. (l <= i < r) && (lower <= v[i] ^ value_xor < upper)
  S range_sum(int l, int r, T lower, T upper, T value_xor = 0) {
    if (lower >= upper) return M::e();
    return range_sum_(B - 1, l, r, T(0), T(1) << B, lower, upper, value_xor);
  }

 private:
  S range_sum_(int h, int l, int r, T vl, T vr, T lower, T upper, T value_xor) {
    if (r <= l) return M::e();
    if (vr <= lower || upper <= vl) return M::e();
    if (lower <= vl && vr <= upper) return seg[h + 1].prod(l, r);

    u32 l0 = bv[h].rank0(l), r0 = bv[h].rank0(r);
    u32 zeros = bv[h].zeros;
    u32 l1 = l + zeros - l0, r1 = r + zeros - r0;
    if ((value_xor >> h) & 1) {
      swap(l0, l1);
      swap(r0, r1);
    }

    T vm = (vl + vr) >> 1;
    return M::op(range_sum_(h - 1, l0, r0, vl, vm, lower, upper, value_xor),
                 range_sum_(h - 1, l1, r1, vm, vr, lower, upper, value_xor));
  }
};
Back to top page