data-structure/wavelet-matrix-with-weight.hpp
Depends on
Verified with
Code
#pragma once
#include "data-structure/wavelet-matrix.hpp"
// W: commutative, inverse
template <class T, class W, int B = 30>
struct WaveletMatrixWithWeight : public WaveletMatrix<T, B> {
using Base = WaveletMatrix<T, B>;
using u32 = uint32_t;
using i64 = int64_t;
using u64 = uint64_t;
using Base::a;
using Base::bv;
using Base::n;
vector<W> w;
vector<vector<W>> sum;
WaveletMatrixWithWeight(u32 _n) : Base(_n), w(_n) {}
WaveletMatrixWithWeight(const vector<T>& _a, const vector<W>& _w) : Base(_a), w(_w) { build(); }
void set(u32 i, const T& x, const W& v) {
assert(x >= 0);
a[i] = x;
w[i] = v;
}
void build() {
bv.assign(B, n);
sum.assign(B + 1, vector<W>(n + 1));
for (int i = 0; i < n; i++) sum[B][i + 1] = sum[B][i] + w[i];
vector<T> cur = a, nxt(n);
vector<W> 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];
}
for (int i = 0; i < n; i++) sum[h][i + 1] = sum[h][i] + wnxt[i];
swap(cur, nxt);
swap(wcur, wnxt);
}
}
// count i s.t. (l <= i < r) && (lower <= v[i] ^ value_xor < upper)
W range_sum(int l, int r, T lower, T upper, T value_xor = 0) {
return range_sum_(l, r, upper, value_xor) - range_sum_(l, r, lower, value_xor);
}
private:
// count i s.t. (l <= i < r) && (v[i] ^ value_xor < upper)
W range_sum_(int l, int r, T upper, T value_xor = 0) {
if (upper >= (T(1) << B)) return sum[B][r] - sum[B][l];
W ret = 0;
for (int h = B - 1; h >= 0; --h) {
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 ((upper >> h) & 1) {
ret += sum[h][r0] - sum[h][l0];
l = l1, r = r1;
} else {
l = l0, r = r0;
}
}
return ret;
}
};
#line 2 "data-structure/wavelet-matrix-with-weight.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 4 "data-structure/wavelet-matrix-with-weight.hpp"
// W: commutative, inverse
template <class T, class W, int B = 30>
struct WaveletMatrixWithWeight : public WaveletMatrix<T, B> {
using Base = WaveletMatrix<T, B>;
using u32 = uint32_t;
using i64 = int64_t;
using u64 = uint64_t;
using Base::a;
using Base::bv;
using Base::n;
vector<W> w;
vector<vector<W>> sum;
WaveletMatrixWithWeight(u32 _n) : Base(_n), w(_n) {}
WaveletMatrixWithWeight(const vector<T>& _a, const vector<W>& _w) : Base(_a), w(_w) { build(); }
void set(u32 i, const T& x, const W& v) {
assert(x >= 0);
a[i] = x;
w[i] = v;
}
void build() {
bv.assign(B, n);
sum.assign(B + 1, vector<W>(n + 1));
for (int i = 0; i < n; i++) sum[B][i + 1] = sum[B][i] + w[i];
vector<T> cur = a, nxt(n);
vector<W> 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];
}
for (int i = 0; i < n; i++) sum[h][i + 1] = sum[h][i] + wnxt[i];
swap(cur, nxt);
swap(wcur, wnxt);
}
}
// count i s.t. (l <= i < r) && (lower <= v[i] ^ value_xor < upper)
W range_sum(int l, int r, T lower, T upper, T value_xor = 0) {
return range_sum_(l, r, upper, value_xor) - range_sum_(l, r, lower, value_xor);
}
private:
// count i s.t. (l <= i < r) && (v[i] ^ value_xor < upper)
W range_sum_(int l, int r, T upper, T value_xor = 0) {
if (upper >= (T(1) << B)) return sum[B][r] - sum[B][l];
W ret = 0;
for (int h = B - 1; h >= 0; --h) {
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 ((upper >> h) & 1) {
ret += sum[h][r0] - sum[h][l0];
l = l1, r = r1;
} else {
l = l0, r = r0;
}
}
return ret;
}
};
Back to top page