ds/wavelet_matrix/wavelet_matrix_2d_range_dynamic_abelgroup.hpp

• View this file on GitHub
• Last update: 2024-04-09 15:17:41+09:00
• Include: #include "ds/wavelet_matrix/wavelet_matrix_2d_range_dynamic_abelgroup.hpp"

Depends on

• alg/monoid/add.hpp
• ds/bit_vector.hpp
• ds/fenwicktree/fenwicktree.hpp

Verified with

• test/library_checker/datastructure/point_add_rectangle_sum_wm_abel.test.cpp
• test/yukicoder/1919_2.test.cpp

Code

#include "ds/bit_vector.hpp"
#include "ds/fenwicktree/fenwicktree.hpp"

template <typename Monoid, typename XY, bool SMALL_X, bool SMALL_Y>
struct Wavelet_Matrix_2D_Range_Dynamic_AbelGroup {
  // 点群を 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;
  // 各段に fenwick tree
  vc<FenwickTree<Monoid>> dat;

  template <typename F>
  Wavelet_Matrix_2D_Range_Dynamic_AbelGroup(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 = __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 count_inner(y1, y2, x2) - count_inner(y1, y2, x1);
  }

  X prod(XY x1, XY x2, XY y1, XY y2) { return sum(x1, x2, y1, y2); }
  X sum(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 add = sum_inner(y1, y2, x2);
    X sub = sum_inner(y1, y2, x1);
    return MX::op(add, MX::inverse(sub));
  }

  X prefix_prod(XY x, XY y) { return prefix_sum(x, y); }
  X prefix_sum(XY x, XY y) {
    if (N == 0) return MX::unit();
    return sum_inner(0, YtoI(y), XtoI(x));
  }

  // 最初に与えた点群の index
  void add(int i, X x) {
    assert(0 <= i && i < N);
    i = new_idx[i];
    int a = A[i];
    FOR_R(d, lg) {
      if (a >> d & 1) {
        i = mid[d] + bv[d].rank(i, 1);
      } else {
        i = bv[d].rank(i, 0);
      }
      dat[d].add(i, x);
    }
  }

private:
  int count_inner(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;
  }

  X sum_inner(int L, int R, int x) {
    if (x == 0) return MX::unit();
    X sm = MX::unit();
    FOR_R(d, lg) {
      int l0 = bv[d].rank(L, 0), r0 = bv[d].rank(R, 0);
      if (x >> d & 1) {
        sm = MX::op(sm, dat[d].sum(l0, r0));
        L += mid[d] - l0, R += mid[d] - r0;
      } else {
        L = l0, R = r0;
      }
    }
    return sm;
  }
};

#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 "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 3 "ds/fenwicktree/fenwicktree.hpp"

template <typename Monoid>
struct FenwickTree {
  using G = Monoid;
  using E = typename G::value_type;
  int n;
  vector<E> dat;
  E total;

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

  void build(int m) {
    n = m;
    dat.assign(m, G::unit());
    total = G::unit();
  }
  void build(const vc<E>& v) {
    build(len(v), [&](int i) -> E { return v[i]; });
  }
  template <typename F>
  void build(int m, F f) {
    n = m;
    dat.clear();
    dat.reserve(n);
    total = G::unit();
    FOR(i, n) { dat.eb(f(i)); }
    for (int i = 1; i <= n; ++i) {
      int j = i + (i & -i);
      if (j <= n) dat[j - 1] = G::op(dat[i - 1], dat[j - 1]);
    }
    total = prefix_sum(m);
  }

  E prod_all() { return total; }
  E sum_all() { return total; }
  E sum(int k) { return prefix_sum(k); }
  E prod(int k) { return prefix_prod(k); }
  E prefix_sum(int k) { return prefix_prod(k); }
  E prefix_prod(int k) {
    chmin(k, n);
    E ret = G::unit();
    for (; k > 0; k -= k & -k) ret = G::op(ret, dat[k - 1]);
    return ret;
  }
  E sum(int L, int R) { return prod(L, R); }
  E prod(int L, int R) {
    chmax(L, 0), chmin(R, n);
    if (L == 0) return prefix_prod(R);
    assert(0 <= L && L <= R && R <= n);
    E pos = G::unit(), neg = G::unit();
    while (L < R) { pos = G::op(pos, dat[R - 1]), R -= R & -R; }
    while (R < L) { neg = G::op(neg, dat[L - 1]), L -= L & -L; }
    return G::op(pos, G::inverse(neg));
  }

  vc<E> get_all() {
    vc<E> res(n);
    FOR(i, n) res[i] = prod(i, i + 1);
    return res;
  }

  void add(int k, E x) { multiply(k, x); }
  void multiply(int k, E x) {
    static_assert(G::commute);
    total = G::op(total, x);
    for (++k; k <= n; k += k & -k) dat[k - 1] = G::op(dat[k - 1], x);
  }

  template <class F>
  int max_right(const F check, int L = 0) {
    assert(check(G::unit()));
    E s = G::unit();
    int i = L;
    // 2^k 進むとダメ
    int k = [&]() {
      while (1) {
        if (i % 2 == 1) { s = G::op(s, G::inverse(dat[i - 1])), i -= 1; }
        if (i == 0) { return topbit(n) + 1; }
        int k = lowbit(i) - 1;
        if (i + (1 << k) > n) return k;
        E t = G::op(s, dat[i + (1 << k) - 1]);
        if (!check(t)) { return k; }
        s = G::op(s, G::inverse(dat[i - 1])), i -= i & -i;
      }
    }();
    while (k) {
      --k;
      if (i + (1 << k) - 1 < len(dat)) {
        E t = G::op(s, dat[i + (1 << k) - 1]);
        if (check(t)) { i += (1 << k), s = t; }
      }
    }
    return i;
  }

  // check(i, x)
  template <class F>
  int max_right_with_index(const F check, int L = 0) {
    assert(check(L, G::unit()));
    E s = G::unit();
    int i = L;
    // 2^k 進むとダメ
    int k = [&]() {
      while (1) {
        if (i % 2 == 1) { s = G::op(s, G::inverse(dat[i - 1])), i -= 1; }
        if (i == 0) { return topbit(n) + 1; }
        int k = lowbit(i) - 1;
        if (i + (1 << k) > n) return k;
        E t = G::op(s, dat[i + (1 << k) - 1]);
        if (!check(i + (1 << k), t)) { return k; }
        s = G::op(s, G::inverse(dat[i - 1])), i -= i & -i;
      }
    }();
    while (k) {
      --k;
      if (i + (1 << k) - 1 < len(dat)) {
        E t = G::op(s, dat[i + (1 << k) - 1]);
        if (check(i + (1 << k), t)) { i += (1 << k), s = t; }
      }
    }
    return i;
  }

  template <class F>
  int min_left(const F check, int R) {
    assert(check(G::unit()));
    E s = G::unit();
    int i = R;
    // false になるところまで戻る
    int k = 0;
    while (i > 0 && check(s)) {
      s = G::op(s, dat[i - 1]);
      k = lowbit(i);
      i -= i & -i;
    }
    if (check(s)) {
      assert(i == 0);
      return 0;
    }
    // 2^k 進むと ok になる
    // false を維持して進む
    while (k) {
      --k;
      E t = G::op(s, G::inverse(dat[i + (1 << k) - 1]));
      if (!check(t)) { i += (1 << k), s = t; }
    }
    return i + 1;
  }

  int kth(E k, int L = 0) {
    return max_right([&k](E x) -> bool { return x <= k; }, L);
  }
};
#line 3 "ds/wavelet_matrix/wavelet_matrix_2d_range_dynamic_abelgroup.hpp"

template <typename Monoid, typename XY, bool SMALL_X, bool SMALL_Y>
struct Wavelet_Matrix_2D_Range_Dynamic_AbelGroup {
  // 点群を 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;
  // 各段に fenwick tree
  vc<FenwickTree<Monoid>> dat;

  template <typename F>
  Wavelet_Matrix_2D_Range_Dynamic_AbelGroup(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 = __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 count_inner(y1, y2, x2) - count_inner(y1, y2, x1);
  }

  X prod(XY x1, XY x2, XY y1, XY y2) { return sum(x1, x2, y1, y2); }
  X sum(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 add = sum_inner(y1, y2, x2);
    X sub = sum_inner(y1, y2, x1);
    return MX::op(add, MX::inverse(sub));
  }

  X prefix_prod(XY x, XY y) { return prefix_sum(x, y); }
  X prefix_sum(XY x, XY y) {
    if (N == 0) return MX::unit();
    return sum_inner(0, YtoI(y), XtoI(x));
  }

  // 最初に与えた点群の index
  void add(int i, X x) {
    assert(0 <= i && i < N);
    i = new_idx[i];
    int a = A[i];
    FOR_R(d, lg) {
      if (a >> d & 1) {
        i = mid[d] + bv[d].rank(i, 1);
      } else {
        i = bv[d].rank(i, 0);
      }
      dat[d].add(i, x);
    }
  }

private:
  int count_inner(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;
  }

  X sum_inner(int L, int R, int x) {
    if (x == 0) return MX::unit();
    X sm = MX::unit();
    FOR_R(d, lg) {
      int l0 = bv[d].rank(L, 0), r0 = bv[d].rank(R, 0);
      if (x >> d & 1) {
        sm = MX::op(sm, dat[d].sum(l0, r0));
        L += mid[d] - l0, R += mid[d] - r0;
      } else {
        L = l0, R = r0;
      }
    }
    return sm;
  }
};

Back to top page