ds/wavelet_matrix/wavelet_matrix_2d_range_static_monoid.hpp

• View this file on GitHub
• Last update: 2024-04-04 05:27:34+09:00
• Include: #include "ds/wavelet_matrix/wavelet_matrix_2d_range_static_monoid.hpp"

Depends on

• alg/monoid/add.hpp
• ds/bit_vector.hpp
• ds/segtree/segtree.hpp
• ds/sparse_table/disjoint_sparse_table.hpp
• ds/sparse_table/sparse_table.hpp
• ds/static_range_product.hpp

Verified with

• test/yukicoder/1600_2.test.cpp

Code

#include "ds/bit_vector.hpp"
#include "ds/segtree/segtree.hpp"
#include "alg/monoid/add.hpp"
#include "ds/static_range_product.hpp"

template <typename Monoid, typename ST, typename XY, bool SMALL_X, bool SMALL_Y>
struct Wavelet_Matrix_2D_Range_Static_Monoid {
  // 点群を Y 昇順に並べる.
  // X を整数になおして binary trie みたいに振り分ける
  using MX = Monoid;
  using X = typename MX::value_type;
  static_assert(MX::commute);

  template <bool SMALL>
  struct TO_IDX {
    vc<XY> key;
    XY mi, ma;
    vc<int> dat;

    void build(vc<XY>& X) {
      if constexpr (SMALL) {
        mi = (X.empty() ? 0 : MIN(X));
        ma = (X.empty() ? 0 : MAX(X));
        dat.assign(ma - mi + 2, 0);
        for (auto& x: X) { dat[x - mi + 1]++; }
        FOR(i, len(dat) - 1) dat[i + 1] += dat[i];
      } else {
        key = X;
        sort(all(key));
      }
    }
    int operator()(XY x) {
      if constexpr (SMALL) {
        return dat[clamp<XY>(x - mi, 0, ma - mi + 1)];
      } else {
        return LB(key, x);
      }
    }
  };

  TO_IDX<SMALL_X> XtoI;
  TO_IDX<SMALL_Y> YtoI;

  int N, lg;
  vector<int> mid;
  vector<Bit_Vector> bv;
  vc<int> new_idx;
  vc<int> A;
  using SEG = Static_Range_Product<Monoid, ST, 4>;
  vc<SEG> dat;

  template <typename F>
  Wavelet_Matrix_2D_Range_Static_Monoid(int N, F f) {
    build(N, f);
  }

  template <typename F>
  void build(int N_, F f) {
    N = N_;
    if (N == 0) {
      lg = 0;
      return;
    }
    vc<XY> tmp(N), Y(N);
    vc<X> S(N);
    FOR(i, N) tie(tmp[i], Y[i], S[i]) = f(i);
    auto I = argsort(Y);
    tmp = rearrange(tmp, I), Y = rearrange(Y, I), S = rearrange(S, I);
    XtoI.build(tmp), YtoI.build(Y);
    new_idx.resize(N);
    FOR(i, N) new_idx[I[i]] = i;

    // あとは普通に
    lg = tmp.empty() ? 0 : __lg(XtoI(MAX(tmp) + 1)) + 1;
    mid.resize(lg), bv.assign(lg, Bit_Vector(N));
    dat.resize(lg);
    A.resize(N);
    FOR(i, N) A[i] = XtoI(tmp[i]);

    vc<XY> A0(N), A1(N);
    vc<X> S0(N), S1(N);
    FOR_R(d, lg) {
      int p0 = 0, p1 = 0;
      FOR(i, N) {
        bool f = (A[i] >> d & 1);
        if (!f) { S0[p0] = S[i], A0[p0] = A[i], p0++; }
        if (f) { S1[p1] = S[i], A1[p1] = A[i], bv[d].set(i), p1++; }
      }
      mid[d] = p0;
      bv[d].build();
      swap(A, A0), swap(S, S0);
      FOR(i, p1) A[p0 + i] = A1[i], S[p0 + i] = S1[i];
      dat[d].build(N, [&](int i) -> X { return S[i]; });
    }
    FOR(i, N) A[i] = XtoI(tmp[i]);
  }

  int count(XY x1, XY x2, XY y1, XY y2) {
    if (N == 0) return 0;
    x1 = XtoI(x1), x2 = XtoI(x2);
    y1 = YtoI(y1), y2 = YtoI(y2);
    return prefix_count(y1, y2, x2) - prefix_count(y1, y2, x1);
  }

  X prod(XY x1, XY x2, XY y1, XY y2) {
    if (N == 0) return MX::unit();
    assert(x1 <= x2 && y1 <= y2);
    x1 = XtoI(x1), x2 = XtoI(x2);
    y1 = YtoI(y1), y2 = YtoI(y2);
    X res = MX::unit();
    prod_dfs(y1, y2, x1, x2, lg - 1, res);
    return res;
  }

private:
  int prefix_count(int L, int R, int x) {
    int cnt = 0;
    FOR_R(d, lg) {
      int l0 = bv[d].rank(L, 0), r0 = bv[d].rank(R, 0);
      if (x >> d & 1) {
        cnt += r0 - l0, L += mid[d] - l0, R += mid[d] - r0;
      } else {
        L = l0, R = r0;
      }
    }
    return cnt;
  }

  void prod_dfs(int L, int R, int x1, int x2, int d, X& res) {
    chmax(x1, 0), chmin(x2, 1 << (d + 1));
    if (x1 >= x2) { return; }
    assert(0 <= x1 && x1 < x2 && x2 <= (1 << (d + 1)));
    if (x1 == 0 && x2 == (1 << (d + 1))) {
      res = MX::op(res, dat[d + 1].prod(L, R));
      return;
    }
    int l0 = bv[d].rank(L, 0), r0 = bv[d].rank(R, 0);
    prod_dfs(l0, r0, x1, x2, d - 1, res);
    prod_dfs(L + mid[d] - l0, R + mid[d] - r0, x1 - (1 << d), x2 - (1 << d),
             d - 1, res);
  }
};

#line 1 "ds/bit_vector.hpp"
struct Bit_Vector {
  vc<pair<u32, u32>> dat;
  Bit_Vector(int n) { dat.assign((n + 63) >> 5, {0, 0}); }

  void set(int i) { dat[i >> 5].fi |= u32(1) << (i & 31); }

  void build() {
    FOR(i, len(dat) - 1) dat[i + 1].se = dat[i].se + popcnt(dat[i].fi);
  }

  // [0, k) 内の 1 の個数
  int rank(int k, bool f = 1) {
    auto [a, b] = dat[k >> 5];
    int ret = b + popcnt(a & ((u32(1) << (k & 31)) - 1));
    return (f ? ret : k - ret);
  }
};
#line 2 "ds/segtree/segtree.hpp"

template <class Monoid>
struct SegTree {
  using MX = Monoid;
  using X = typename MX::value_type;
  using value_type = X;
  vc<X> dat;
  int n, log, size;

  SegTree() {}
  SegTree(int n) { build(n); }
  template <typename F>
  SegTree(int n, F f) {
    build(n, f);
  }
  SegTree(const vc<X>& v) { build(v); }

  void build(int m) {
    build(m, [](int i) -> X { return MX::unit(); });
  }
  void build(const vc<X>& v) {
    build(len(v), [&](int i) -> X { return v[i]; });
  }
  template <typename F>
  void build(int m, F f) {
    n = m, log = 1;
    while ((1 << log) < n) ++log;
    size = 1 << log;
    dat.assign(size << 1, MX::unit());
    FOR(i, n) dat[size + i] = f(i);
    FOR_R(i, 1, size) update(i);
  }

  X get(int i) { return dat[size + i]; }
  vc<X> get_all() { return {dat.begin() + size, dat.begin() + size + n}; }

  void update(int i) { dat[i] = Monoid::op(dat[2 * i], dat[2 * i + 1]); }
  void set(int i, const X& x) {
    assert(i < n);
    dat[i += size] = x;
    while (i >>= 1) update(i);
  }

  void multiply(int i, const X& x) {
    assert(i < n);
    i += size;
    dat[i] = Monoid::op(dat[i], x);
    while (i >>= 1) update(i);
  }

  X prod(int L, int R) {
    assert(0 <= L && L <= R && R <= n);
    X vl = Monoid::unit(), vr = Monoid::unit();
    L += size, R += size;
    while (L < R) {
      if (L & 1) vl = Monoid::op(vl, dat[L++]);
      if (R & 1) vr = Monoid::op(dat[--R], vr);
      L >>= 1, R >>= 1;
    }
    return Monoid::op(vl, vr);
  }

  X prod_all() { return dat[1]; }

  template <class F>
  int max_right(F check, int L) {
    assert(0 <= L && L <= n && check(Monoid::unit()));
    if (L == n) return n;
    L += size;
    X sm = Monoid::unit();
    do {
      while (L % 2 == 0) L >>= 1;
      if (!check(Monoid::op(sm, dat[L]))) {
        while (L < size) {
          L = 2 * L;
          if (check(Monoid::op(sm, dat[L]))) { sm = Monoid::op(sm, dat[L++]); }
        }
        return L - size;
      }
      sm = Monoid::op(sm, dat[L++]);
    } while ((L & -L) != L);
    return n;
  }

  template <class F>
  int min_left(F check, int R) {
    assert(0 <= R && R <= n && check(Monoid::unit()));
    if (R == 0) return 0;
    R += size;
    X sm = Monoid::unit();
    do {
      --R;
      while (R > 1 && (R % 2)) R >>= 1;
      if (!check(Monoid::op(dat[R], sm))) {
        while (R < size) {
          R = 2 * R + 1;
          if (check(Monoid::op(dat[R], sm))) { sm = Monoid::op(dat[R--], sm); }
        }
        return R + 1 - size;
      }
      sm = Monoid::op(dat[R], sm);
    } while ((R & -R) != R);
    return 0;
  }

  // prod_{l<=i<r} A[i xor x]
  X xor_prod(int l, int r, int xor_val) {
    static_assert(Monoid::commute);
    X x = Monoid::unit();
    for (int k = 0; k < log + 1; ++k) {
      if (l >= r) break;
      if (l & 1) { x = Monoid::op(x, dat[(size >> k) + ((l++) ^ xor_val)]); }
      if (r & 1) { x = Monoid::op(x, dat[(size >> k) + ((--r) ^ xor_val)]); }
      l /= 2, r /= 2, xor_val /= 2;
    }
    return x;
  }
};
#line 2 "alg/monoid/add.hpp"

template <typename E>
struct Monoid_Add {
  using X = E;
  using value_type = X;
  static constexpr X op(const X &x, const X &y) noexcept { return x + y; }
  static constexpr X inverse(const X &x) noexcept { return -x; }
  static constexpr X power(const X &x, ll n) noexcept { return X(n) * x; }
  static constexpr X unit() { return X(0); }
  static constexpr bool commute = true;
};
#line 2 "ds/sparse_table/sparse_table.hpp"

// 冪等なモノイドであることを仮定。disjoint sparse table より x 倍高速
template <class Monoid>
struct Sparse_Table {
  using MX = Monoid;
  using X = typename MX::value_type;
  int n, log;
  vvc<X> dat;

  Sparse_Table() {}
  Sparse_Table(int n) { build(n); }
  template <typename F>
  Sparse_Table(int n, F f) {
    build(n, f);
  }
  Sparse_Table(const vc<X>& v) { build(v); }

  void build(int m) {
    build(m, [](int i) -> X { return MX::unit(); });
  }
  void build(const vc<X>& v) {
    build(len(v), [&](int i) -> X { return v[i]; });
  }
  template <typename F>
  void build(int m, F f) {
    n = m, log = 1;
    while ((1 << log) < n) ++log;
    dat.resize(log);
    dat[0].resize(n);
    FOR(i, n) dat[0][i] = f(i);

    FOR(i, log - 1) {
      dat[i + 1].resize(len(dat[i]) - (1 << i));
      FOR(j, len(dat[i]) - (1 << i)) {
        dat[i + 1][j] = MX::op(dat[i][j], dat[i][j + (1 << i)]);
      }
    }
  }

  X prod(int L, int R) {
    if (L == R) return MX::unit();
    if (R == L + 1) return dat[0][L];
    int k = topbit(R - L - 1);
    return MX::op(dat[k][L], dat[k][R - (1 << k)]);
  }

  template <class F>
  int max_right(const F check, int L) {
    assert(0 <= L && L <= n && check(MX::unit()));
    if (L == n) return n;
    int ok = L, ng = n + 1;
    while (ok + 1 < ng) {
      int k = (ok + ng) / 2;
      bool bl = check(prod(L, k));
      if (bl) ok = k;
      if (!bl) ng = k;
    }
    return ok;
  }

  template <class F>
  int min_left(const F check, int R) {
    assert(0 <= R && R <= n && check(MX::unit()));
    if (R == 0) return 0;
    int ok = R, ng = -1;
    while (ng + 1 < ok) {
      int k = (ok + ng) / 2;
      bool bl = check(prod(k, R));
      if (bl) ok = k;
      if (!bl) ng = k;
    }
    return ok;
  }
};
#line 2 "ds/sparse_table/disjoint_sparse_table.hpp"

template <class Monoid>
struct Disjoint_Sparse_Table {
  using MX = Monoid;
  using X = typename MX::value_type;
  int n, log;
  vvc<X> dat;

  Disjoint_Sparse_Table() {}
  Disjoint_Sparse_Table(int n) { build(n); }
  template <typename F>
  Disjoint_Sparse_Table(int n, F f) {
    build(n, f);
  }
  Disjoint_Sparse_Table(const vc<X>& v) { build(v); }

  void build(int m) {
    build(m, [](int i) -> X { return MX::unit(); });
  }
  void build(const vc<X>& v) {
    build(len(v), [&](int i) -> X { return v[i]; });
  }
  template <typename F>
  void build(int m, F f) {
    n = m, log = 1;
    while ((1 << log) < n) ++log;
    dat.resize(log);
    dat[0].reserve(n);
    FOR(i, n) dat[0].eb(f(i));
    FOR(i, 1, log) {
      auto& v = dat[i];
      v = dat[0];
      int b = 1 << i;
      for (int m = b; m <= n; m += 2 * b) {
        int L = m - b, R = min(n, m + b);
        FOR_R(j, L + 1, m) v[j - 1] = MX::op(v[j - 1], v[j]);
        FOR(j, m, R - 1) v[j + 1] = MX::op(v[j], v[j + 1]);
      }
    }
  }

  X prod(int L, int R) {
    if (L == R) return MX::unit();
    --R;
    if (L == R) return dat[0][L];
    int k = topbit(L ^ R);
    return MX::op(dat[k][L], dat[k][R]);
  }

  template <class F>
  int max_right(const F check, int L) {
    assert(0 <= L && L <= n && check(MX::unit()));
    if (L == n) return n;
    int ok = L, ng = n + 1;
    while (ok + 1 < ng) {
      int k = (ok + ng) / 2;
      bool bl = check(prod(L, k));
      if (bl) ok = k;
      if (!bl) ng = k;
    }
    return ok;
  }

  template <class F>
  int min_left(const F check, int R) {
    assert(0 <= R && R <= n && check(MX::unit()));
    if (R == 0) return 0;
    int ok = R, ng = -1;
    while (ng + 1 < ok) {
      int k = (ok + ng) / 2;
      bool bl = check(prod(k, R));
      if (bl) ok = k;
      if (!bl) ng = k;
    }
    return ok;
  }
};
#line 3 "ds/static_range_product.hpp"

/*
参考：https://judge.yosupo.jp/submission/106668
長さ 2^LOG のブロックに分ける．ブロック内の prefix, suffix を持つ．
ブロック積の列を ST(DST) で持つ．ブロックをまたぐ積は O(1).
短いものは O(1) を諦めて愚直ということにする．
前計算：O(Nlog(N)/2^LOG)
クエリ：O(1) / worst O(2^LOG)
*/
template <typename Monoid, typename SPARSE_TABLE, int LOG = 4>
struct Static_Range_Product {
  using MX = Monoid;
  using X = typename MX::value_type;
  int N, b_num;
  vc<X> A, pre, suf; // inclusive
  SPARSE_TABLE ST;

  Static_Range_Product() {}
  template <typename F>
  Static_Range_Product(int n, F f) {
    build(n, f);
  }
  Static_Range_Product(const vc<X>& v) { build(v); }

  void build(const vc<X>& v) {
    build(len(v), [&](int i) -> X { return v[i]; });
  }
  template <typename F>
  void build(int m, F f) {
    N = m;
    b_num = N >> LOG;
    A.resize(N);
    FOR(i, N) A[i] = f(i);
    pre = A, suf = A;
    constexpr int mask = (1 << LOG) - 1;
    FOR(i, 1, N) {
      if (i & mask) pre[i] = MX::op(pre[i - 1], A[i]);
    }
    FOR_R(i, 1, N) {
      if (i & mask) suf[i - 1] = MX::op(A[i - 1], suf[i]);
    }
    ST.build(b_num, [&](int i) -> X { return suf[i << LOG]; });
  }

  // O(1) or O(R-L)
  X prod(int L, int R) {
    if (L == R) return MX::unit();
    R -= 1;
    int a = L >> LOG, b = R >> LOG;
    if (a < b) {
      X x = ST.prod(a + 1, b);
      x = MX::op(suf[L], x);
      x = MX::op(x, pre[R]);
      return x;
    }
    X x = A[L];
    FOR(i, L + 1, R + 1) x = MX::op(x, A[i]);
    return x;
  }
};
#line 5 "ds/wavelet_matrix/wavelet_matrix_2d_range_static_monoid.hpp"

template <typename Monoid, typename ST, typename XY, bool SMALL_X, bool SMALL_Y>
struct Wavelet_Matrix_2D_Range_Static_Monoid {
  // 点群を Y 昇順に並べる.
  // X を整数になおして binary trie みたいに振り分ける
  using MX = Monoid;
  using X = typename MX::value_type;
  static_assert(MX::commute);

  template <bool SMALL>
  struct TO_IDX {
    vc<XY> key;
    XY mi, ma;
    vc<int> dat;

    void build(vc<XY>& X) {
      if constexpr (SMALL) {
        mi = (X.empty() ? 0 : MIN(X));
        ma = (X.empty() ? 0 : MAX(X));
        dat.assign(ma - mi + 2, 0);
        for (auto& x: X) { dat[x - mi + 1]++; }
        FOR(i, len(dat) - 1) dat[i + 1] += dat[i];
      } else {
        key = X;
        sort(all(key));
      }
    }
    int operator()(XY x) {
      if constexpr (SMALL) {
        return dat[clamp<XY>(x - mi, 0, ma - mi + 1)];
      } else {
        return LB(key, x);
      }
    }
  };

  TO_IDX<SMALL_X> XtoI;
  TO_IDX<SMALL_Y> YtoI;

  int N, lg;
  vector<int> mid;
  vector<Bit_Vector> bv;
  vc<int> new_idx;
  vc<int> A;
  using SEG = Static_Range_Product<Monoid, ST, 4>;
  vc<SEG> dat;

  template <typename F>
  Wavelet_Matrix_2D_Range_Static_Monoid(int N, F f) {
    build(N, f);
  }

  template <typename F>
  void build(int N_, F f) {
    N = N_;
    if (N == 0) {
      lg = 0;
      return;
    }
    vc<XY> tmp(N), Y(N);
    vc<X> S(N);
    FOR(i, N) tie(tmp[i], Y[i], S[i]) = f(i);
    auto I = argsort(Y);
    tmp = rearrange(tmp, I), Y = rearrange(Y, I), S = rearrange(S, I);
    XtoI.build(tmp), YtoI.build(Y);
    new_idx.resize(N);
    FOR(i, N) new_idx[I[i]] = i;

    // あとは普通に
    lg = tmp.empty() ? 0 : __lg(XtoI(MAX(tmp) + 1)) + 1;
    mid.resize(lg), bv.assign(lg, Bit_Vector(N));
    dat.resize(lg);
    A.resize(N);
    FOR(i, N) A[i] = XtoI(tmp[i]);

    vc<XY> A0(N), A1(N);
    vc<X> S0(N), S1(N);
    FOR_R(d, lg) {
      int p0 = 0, p1 = 0;
      FOR(i, N) {
        bool f = (A[i] >> d & 1);
        if (!f) { S0[p0] = S[i], A0[p0] = A[i], p0++; }
        if (f) { S1[p1] = S[i], A1[p1] = A[i], bv[d].set(i), p1++; }
      }
      mid[d] = p0;
      bv[d].build();
      swap(A, A0), swap(S, S0);
      FOR(i, p1) A[p0 + i] = A1[i], S[p0 + i] = S1[i];
      dat[d].build(N, [&](int i) -> X { return S[i]; });
    }
    FOR(i, N) A[i] = XtoI(tmp[i]);
  }

  int count(XY x1, XY x2, XY y1, XY y2) {
    if (N == 0) return 0;
    x1 = XtoI(x1), x2 = XtoI(x2);
    y1 = YtoI(y1), y2 = YtoI(y2);
    return prefix_count(y1, y2, x2) - prefix_count(y1, y2, x1);
  }

  X prod(XY x1, XY x2, XY y1, XY y2) {
    if (N == 0) return MX::unit();
    assert(x1 <= x2 && y1 <= y2);
    x1 = XtoI(x1), x2 = XtoI(x2);
    y1 = YtoI(y1), y2 = YtoI(y2);
    X res = MX::unit();
    prod_dfs(y1, y2, x1, x2, lg - 1, res);
    return res;
  }

private:
  int prefix_count(int L, int R, int x) {
    int cnt = 0;
    FOR_R(d, lg) {
      int l0 = bv[d].rank(L, 0), r0 = bv[d].rank(R, 0);
      if (x >> d & 1) {
        cnt += r0 - l0, L += mid[d] - l0, R += mid[d] - r0;
      } else {
        L = l0, R = r0;
      }
    }
    return cnt;
  }

  void prod_dfs(int L, int R, int x1, int x2, int d, X& res) {
    chmax(x1, 0), chmin(x2, 1 << (d + 1));
    if (x1 >= x2) { return; }
    assert(0 <= x1 && x1 < x2 && x2 <= (1 << (d + 1)));
    if (x1 == 0 && x2 == (1 << (d + 1))) {
      res = MX::op(res, dat[d + 1].prod(L, R));
      return;
    }
    int l0 = bv[d].rank(L, 0), r0 = bv[d].rank(R, 0);
    prod_dfs(l0, r0, x1, x2, d - 1, res);
    prod_dfs(L + mid[d] - l0, R + mid[d] - r0, x1 - (1 << d), x2 - (1 << d),
             d - 1, res);
  }
};

Back to top page