競技プログラミング用 C++ ライブラリ
data_structure/point_set_wavelet_matrix.hpp · GitHub で見る
共通ページ
#include "lib/data_structure/point_set_wavelet_matrix.hpp"
std::vector<int> a{3, 1, 4, 1, 5};
std::vector<std::pair<int, int>> candidates{ {0, 10}, {2, 7} };
PointSetWaveletMatrix<int> wm(a, candidates);
wm.set(0, 10);
auto answer = wm.min_count_largest(0, 5, 12);
配列の一点代入と、区間内の順位・和のクエリを扱います。初期値以外に代入する全ての (位置,値) の事前登録が必要です。通常は構築前に更新クエリを先読みし、候補を candidates に集めます。初期値は自動登録されます。候補が最初から分かっていれば、更新クエリ自体の先読みは不要です。順位・和を求めるクエリは先読みせずに処理できます。区間は \([l,r)\)、順位は 0 始まりです。
一点操作と候補の位置は 0 <= i < size()、区間は 0 <= l <= r <= size() とします。kth_* は 0 <= k < r-l、sum_smallest・sum_largest は 0 <= k <= r-l が必要です。和の k は個数で、0 個の和は 0 です。重複する要素は個数分だけ数え、空区間の個数・和は 0 です。同じ候補を複数登録しても一つにまとめます。列長と異なる登録 (位置,値) の総数はそれぞれ INT_MAX 以下とします。入力は変更せず、構築後に破棄できます。
| 操作 | 内容・返り値 | 計算量 |
|---|---|---|
PointSetWaveletMatrix() |
空で構築 | \(O(1)\) |
PointSetWaveletMatrix(a,candidates = {}) / build(a,candidates = {})
|
初期配列と代入候補から構築/再構築 | \(O(m\log m)\) |
size() / empty()
|
配列の長さ/空判定 | \(O(1)\) |
access(i) / operator[](i)
|
現在の値 | \(O(1)\) |
contains(i,x) |
位置 i に値 x を代入できるか |
\(O(\log m)\) |
set(i,x) |
一点代入。未登録なら false を返し、内容を保持 | \(O(\log m\log\sigma)\) |
count(l,r,x) / rank(x,r)
|
区間/\([0,r)\) の x の個数 |
\(O(\log m\log\sigma)\) |
range_freq(l,r,upper) |
upper 未満の個数 |
\(O(\log m\log\sigma)\) |
range_freq(l,r,lower,upper) |
lower 以上 upper 未満の個数 |
\(O(\log m\log\sigma)\) |
range_sum(l,r) |
区間和 | \(O(\log m)\) |
range_sum(l,r,upper) |
upper 未満の値の和 |
\(O(\log m\log\sigma)\) |
range_sum(l,r,lower,upper) |
指定した値域の和 | \(O(\log m\log\sigma)\) |
kth_smallest(l,r,k) / kth_largest(l,r,k)
|
小さい順/大きい順の k 番目 |
\(O(\log m\log\sigma)\) |
sum_smallest(l,r,k) / sum_largest(l,r,k)
|
下位/上位 k 個の和 |
\(O(\log m\log\sigma)\) |
min_count_smallest(l,r,target) / min_count_largest(l,r,target)
|
各順序で和が target 以上になる最小個数。届かなければ -1 |
\(O(\log m\log\sigma)\) |
max_count_smallest(l,r,budget) / max_count_largest(l,r,budget)
|
各順序で和が budget 以下の最大個数。空の選択も不可能なら -1 |
\(O(\log m\log\sigma)\) |
T は bool を除く 64 bit 以下の整数型、Sum は符号付き整数型です。PointSetWaveletMatrix<T,Sum> で和の型を指定できます。既定は 32 bit 以下の値なら std::int64_t、64 bit の値なら __int128_t です。中間の部分和も Sum に収めてください。
和の境界探索 min_count_*・max_count_* には初期値と登録した候補全てが非負であることが必要です。前者は target <= 0 なら 0、後者は budget < 0 なら -1 を返します。その他のクエリは負の値も扱えます。
異なる登録 (位置,値) の総数を \(m\)、異なる値の数を \(\sigma\) として、構築は \(O(m\log m)\) 時間、空間は \(O(m\log\sigma)\) です。更新と区間の順位・和の探索は \(O(\log m\log\sigma)\)、値域を指定しない区間和は \(O(\log m)\)、候補の確認は \(O(\log m)\)、サイズと要素取得は \(O(1)\) です。
#ifndef POINT_SET_WAVELET_MATRIX_HPP
#define POINT_SET_WAVELET_MATRIX_HPP
#include "weighted_wavelet_matrix.hpp"
template <class T = std::int64_t, class Sum = std::conditional_t<(sizeof(T) < 8), std::int64_t, __int128_t>>
class PointSetWaveletMatrix {
std::vector<T> keys;
std::vector<int> offsets{0}, current;
WeightedWaveletMatrix<T, int, Sum> wm;
void check([[maybe_unused]] int l, [[maybe_unused]] int r) const { assert(0 <= l && l <= r && r <= size()); }
int find(int p, const T &x) const {
assert(0 <= p && p < size());
return std::lower_bound(keys.begin() + offsets[p], keys.begin() + offsets[p + 1], x) - keys.begin();
}
public:
using value_type = T;
using sum_type = Sum;
PointSetWaveletMatrix() = default;
explicit PointSetWaveletMatrix(const std::vector<T> &a, const std::vector<std::pair<int, T>> &candidates = {}) { build(a, candidates); }
void build(const std::vector<T> &a, const std::vector<std::pair<int, T>> &candidates = {}) {
assert(a.size() <= std::size_t(std::numeric_limits<int>::max()));
std::vector<std::pair<int, T>> points = candidates;
points.reserve(a.size() + candidates.size());
for (int i = 0; i < static_cast<int>(a.size()); ++i) points.emplace_back(i, a[i]);
std::sort(points.begin(), points.end());
points.erase(std::unique(points.begin(), points.end()), points.end());
assert(points.size() <= std::size_t(std::numeric_limits<int>::max()));
offsets.assign(a.size() + 1, 0);
keys.clear();
keys.reserve(points.size());
for (const auto &[p, x] : points) {
assert(0 <= p && p < static_cast<int>(a.size()));
++offsets[p + 1];
keys.push_back(x);
}
for (int i = 0; i < static_cast<int>(a.size()); ++i) offsets[i + 1] += offsets[i];
current.resize(a.size());
std::vector<int> weights(points.size());
for (int i = 0; i < size(); ++i) weights[current[i] = find(i, a[i])] = 1;
wm.build(keys, weights);
}
int size() const { return static_cast<int>(current.size()); }
bool empty() const { return current.empty(); }
T access(int p) const { assert(0 <= p && p < size()); return keys[current[p]]; }
T operator[](int p) const { return access(p); }
bool contains(int p, const T &x) const {
int q = find(p, x);
return q < offsets[p + 1] && keys[q] == x;
}
bool set(int p, const T &x) {
int q = find(p, x);
if (q == offsets[p + 1] || keys[q] != x) return false;
wm.move_weight(current[p], q, 1);
current[p] = q;
return true;
}
int count(int l, int r, const T &x) const { check(l, r); return wm.count(offsets[l], offsets[r], x); }
int rank(const T &x, int r) const { return count(0, r, x); }
int range_freq(int l, int r, const T &upper) const { check(l, r); return wm.range_freq(offsets[l], offsets[r], upper); }
int range_freq(int l, int r, const T &lower, const T &upper) const { check(l, r); return wm.range_freq(offsets[l], offsets[r], lower, upper); }
Sum range_sum(int l, int r) const { check(l, r); return wm.range_sum(offsets[l], offsets[r]); }
Sum range_sum(int l, int r, const T &upper) const { check(l, r); return wm.range_sum(offsets[l], offsets[r], upper); }
Sum range_sum(int l, int r, const T &lower, const T &upper) const { check(l, r); return wm.range_sum(offsets[l], offsets[r], lower, upper); }
T kth_smallest(int l, int r, int k) const { check(l, r); return wm.kth_smallest(offsets[l], offsets[r], k); }
T kth_largest(int l, int r, int k) const { check(l, r); return wm.kth_largest(offsets[l], offsets[r], k); }
Sum sum_smallest(int l, int r, int k) const { check(l, r); return wm.sum_smallest(offsets[l], offsets[r], k); }
Sum sum_largest(int l, int r, int k) const { check(l, r); return wm.sum_largest(offsets[l], offsets[r], k); }
int min_count_smallest(int l, int r, Sum target) const { check(l, r); return wm.min_count_smallest(offsets[l], offsets[r], target); }
int min_count_largest(int l, int r, Sum target) const { check(l, r); return wm.min_count_largest(offsets[l], offsets[r], target); }
int max_count_smallest(int l, int r, Sum budget) const { check(l, r); return wm.max_count_smallest(offsets[l], offsets[r], budget); }
int max_count_largest(int l, int r, Sum budget) const { check(l, r); return wm.max_count_largest(offsets[l], offsets[r], budget); }
};
#endif
#line 1 "data_structure/point_set_wavelet_matrix.hpp"
#line 1 "data_structure/weighted_wavelet_matrix.hpp"
#line 1 "components/wavelet_compression.hpp"
#include <algorithm>
#include <array>
#include <bit>
#include <cassert>
#include <cstdint>
#include <limits>
#include <type_traits>
#include <vector>
namespace wavelet_detail {
template<class T>
std::uint32_t lower_code(const std::vector<T> &values, bool dense, const T &x) {
if (values.empty()) return 0;
if constexpr (std::is_integral_v<T> && sizeof(T) <= 8)
if (dense) {
if (x <= values.front()) return 0;
if (values.back() < x) return static_cast<int>(values.size());
return std::uint64_t(x) - std::uint64_t(values.front());
}
int p = 0, length = static_cast<int>(values.size());
while (length > 1) {
int half = length / 2;
p += (values[p + half] < x) * half;
length -= half;
}
return p + (values[p] < x);
}
template<class T, bool WithPositions = true>
std::vector<std::uint32_t> compress(const std::vector<T> &a, std::vector<T> &values,
int &ordered, std::vector<int> &positions, std::vector<int> &offsets) {
assert(a.size() <= std::size_t(std::numeric_limits<int>::max()));
const int n = static_cast<int>(a.size());
ordered = 0;
if constexpr (WithPositions) offsets.clear();
std::vector<std::uint32_t> current(n);
bool ascending = std::is_sorted(a.begin(), a.end());
if (ascending || std::is_sorted(a.rbegin(), a.rend())) {
ordered = ascending ? 1 : -1;
values.clear();
values.reserve(n);
for (int i = 0; i < n; ++i) {
int j = ascending ? i : n - 1 - i;
if (values.empty() || values.back() != a[j]) values.push_back(a[j]);
current[j] = static_cast<int>(values.size()) - 1;
}
if (values.capacity() > 2 * values.size()) values.shrink_to_fit();
if constexpr (WithPositions) positions.clear();
return current;
}
if constexpr (std::is_integral_v<T> && !std::is_same_v<T, bool> && sizeof(T) <= 8) {
if (n >= 1024) {
using U = std::make_unsigned_t<T>;
struct [[gnu::packed]] Item {
U key;
std::uint32_t index;
};
std::vector<Item> order(n);
U lo = std::numeric_limits<U>::max(), hi = 0;
constexpr U sign = std::is_signed_v<T> ? U(1) << (sizeof(T) * 8 - 1) : U(0);
for (int i = 0; i < n; ++i) {
U key = U(a[i]) ^ sign;
order[i] = {key, std::uint32_t(i)};
lo = std::min(lo, key);
hi = std::max(hi, key);
}
if (hi - lo < 65536 && std::uint64_t(hi - lo) < unsigned(n)) {
std::vector<int> codes(std::size_t(hi - lo) + 1, -1);
for (const auto &v : order) codes[v.key - lo] = 0;
values.clear();
values.reserve(codes.size());
for (std::size_t i = 0; i < codes.size(); ++i)
if (codes[i] == 0) {
codes[i] = static_cast<int>(values.size());
values.push_back(T(U(lo + U(i)) ^ sign));
}
for (const auto &v : order) current[v.index] = codes[v.key - lo];
if (values.capacity() > 2 * values.size()) values.shrink_to_fit();
if constexpr (WithPositions) positions.clear();
return current;
}
std::vector<Item> next(n);
int width = std::bit_width(U(lo ^ hi)), span_width = std::bit_width(U(hi - lo));
if ((width + 10) / 11 > (span_width + 10) / 11) {
for (auto &v : order) v.key -= lo;
width = span_width;
}
for (int shift = 0; shift < width; shift += 11) {
std::array<int, 2048> offsets{};
for (const auto &v : order) ++offsets[(v.key >> shift) & 2047];
int total = 0;
for (int &c : offsets) {
int t = c;
c = total;
total += t;
}
for (const auto &v : order) next[offsets[(v.key >> shift) & 2047]++] = v;
order.swap(next);
}
values.clear();
values.reserve(n);
U previous = 0;
if constexpr (WithPositions) {
positions.resize(n);
offsets.reserve(std::size_t(n) + 1);
}
for (int j = 0; j < n; ++j) {
const auto &v = order[j];
if (values.empty() || v.key != previous) {
values.push_back(a[v.index]);
if constexpr (WithPositions) offsets.push_back(j);
}
previous = v.key;
current[v.index] = static_cast<int>(values.size()) - 1;
if constexpr (WithPositions) positions[j] = v.index;
}
if constexpr (WithPositions) offsets.push_back(n);
if (values.capacity() > 2 * values.size()) values.shrink_to_fit();
return current;
}
}
values = a;
std::sort(values.begin(), values.end());
values.erase(std::unique(values.begin(), values.end()), values.end());
if (values.capacity() > 2 * values.size()) values.shrink_to_fit();
for (int i = 0; i < n; ++i) current[i] = lower_code<T>(values, false, a[i]);
if constexpr (WithPositions) positions.clear();
return current;
}
template<class T> struct ValueIndex {
int n = 0;
std::vector<T> values;
std::vector<std::uint32_t> original;
void build(const std::vector<T> &a) {
assert(a.size() <= std::size_t(std::numeric_limits<int>::max()));
n = static_cast<int>(a.size());
int ordered = 0;
std::vector<int> positions, offsets;
original = compress<T, false>(a, values, ordered, positions, offsets);
}
std::uint32_t code(const T &x) const { return lower_code(values, false, x); }
};
}
#line 1 "components/quad_bit_vector.hpp"
#line 8 "components/quad_bit_vector.hpp"
#include <cstddef>
#line 11 "components/quad_bit_vector.hpp"
#if defined(__x86_64__) || defined(__i386__)
#include <immintrin.h>
#endif
struct QuadBitVector {
struct Block {
std::uint64_t low = 0, high = 0;
std::array<int, 4> prefix{};
};
int n;
std::vector<Block> blocks;
#if defined(__x86_64__) || defined(__i386__)
__attribute__((target("avx2"))) static void pack_avx2(const std::uint32_t *a, int n, int shift, Block *blocks) {
auto low_shift = _mm_cvtsi32_si128(31 - shift), high_shift = _mm_cvtsi32_si128(30 - shift);
for (int w = 0; w < (n >> 6) + bool(n & 63); ++w) {
int start = w * 64;
std::uint64_t low = 0, high = 0;
int length = std::min(64, n - start), j = 0;
for (; j + 8 <= length; j += 8) {
auto x = _mm256_loadu_si256(reinterpret_cast<const __m256i *>(a + start + j));
low |= std::uint64_t(_mm256_movemask_ps(_mm256_castsi256_ps(_mm256_sll_epi32(x, low_shift)))) << j;
high |= std::uint64_t(_mm256_movemask_ps(_mm256_castsi256_ps(_mm256_sll_epi32(x, high_shift)))) << j;
}
for (; j < length; ++j) {
low |= std::uint64_t((a[start + j] >> shift) & 1) << j;
high |= std::uint64_t((a[start + j] >> (shift + 1)) & 1) << j;
}
blocks[w].low = low;
blocks[w].high = high;
}
blocks[n >> 6].low = n & 63 ? blocks[n >> 6].low : 0;
blocks[n >> 6].high = n & 63 ? blocks[n >> 6].high : 0;
}
#endif
explicit QuadBitVector(int n = 0) : n(n), blocks((n >> 6) + 1) {}
void build(const std::vector<std::uint32_t> &a, int shift) {
assert(a.size() == std::size_t(n) && 0 <= shift && shift <= 30);
[[maybe_unused]] bool vectorized = false;
#if defined(__x86_64__) || defined(__i386__)
vectorized = n >= 64 && __builtin_cpu_supports("avx2");
if (vectorized) pack_avx2(a.data(), n, shift, blocks.data());
#endif
std::array<int, 4> counts{};
for (int w = 0; w < static_cast<int>(blocks.size()); ++w) {
std::uint64_t low = 0, high = 0;
int start = w * 64, length = std::min(64, n - start);
if (vectorized) low = blocks[w].low, high = blocks[w].high;
else
for (int j = 0; j < length; ++j) {
low |= std::uint64_t((a[start + j] >> shift) & 1) << j;
high |= std::uint64_t((a[start + j] >> (shift + 1)) & 1) << j;
}
blocks[w] = {low, high, counts};
int c1 = std::popcount(low & ~high), c2 = std::popcount(~low & high), c3 = std::popcount(low & high);
counts[0] += length - c1 - c2 - c3;
counts[1] += c1;
counts[2] += c2;
counts[3] += c3;
}
}
std::uint64_t mask(const Block &b, int digit) const {
return ((digit & 1) ? b.low : ~b.low) & ((digit & 2) ? b.high : ~b.high);
}
int rank(int r, int digit) const {
const auto &b = blocks[r >> 6];
return b.prefix[digit] + std::popcount(mask(b, digit) & ((std::uint64_t(1) << (r & 63)) - 1));
}
int below(int r, int digit) const {
if (!digit) return 0;
if (digit == 1) return rank(r, 0);
if (digit == 3) return r - rank(r, 3);
const auto &b = blocks[r >> 6];
return b.prefix[0] + b.prefix[1] + std::popcount(~b.high & ((std::uint64_t(1) << (r & 63)) - 1));
}
std::array<int, 4> ranks(int r) const {
const auto &b = blocks[r >> 6];
std::uint64_t m = (std::uint64_t(1) << (r & 63)) - 1;
std::array<int, 4> a;
a[1] = b.prefix[1] + std::popcount(b.low & ~b.high & m);
a[2] = b.prefix[2] + std::popcount(~b.low & b.high & m);
a[3] = b.prefix[3] + std::popcount(b.low & b.high & m);
a[0] = r - a[1] - a[2] - a[3];
return a;
}
};
#line 1 "components/fenwick_sum.hpp"
#line 8 "components/fenwick_sum.hpp"
template <class T>
class FenwickSum {
std::vector<T> data;
public:
FenwickSum() = default;
explicit FenwickSum(int n) { assert(n >= 0); data.resize(n); }
explicit FenwickSum(const std::vector<T> &a) { build(a); }
void build(const std::vector<T> &a) {
assert(a.size() <= static_cast<unsigned>(std::numeric_limits<int>::max()));
data = a;
constexpr unsigned step = sizeof(T) <= 8 ? 2 : 1;
for (unsigned i = step; i <= data.size(); i += step) {
if constexpr (step == 2) data[i - 1] += data[i - 2];
unsigned p = i + (i & -i);
if (p <= data.size()) data[p - 1] += data[i - 1];
}
}
int size() const { return static_cast<int>(data.size()); }
void add(int p, const T &x) {
assert(0 <= p && p < size());
for (unsigned i = unsigned(p) + 1; i <= data.size(); i += i & -i) data[i - 1] += x;
}
void add_pair(int p, const T &x, int q, const T &y) {
assert(0 <= p && p < size() && 0 <= q && q < size());
unsigned a = unsigned(p) + 1, b = unsigned(q) + 1;
while (a != b && (a <= data.size() || b <= data.size())) {
if (a < b) data[a - 1] += x, a += a & -a;
else data[b - 1] += y, b += b & -b;
}
T both = x;
both += y;
for (; a <= data.size(); a += a & -a) data[a - 1] += both;
}
T prefix(int r) const {
assert(0 <= r && r <= size());
T result{};
for (; r; r &= r - 1) result += data[r - 1];
return result;
}
template <class Accum>
Accum prod_as(int l, int r) const {
assert(0 <= l && l <= r && r <= size());
Accum result{};
if (l == r) return result;
unsigned shared = unsigned(l) & (std::numeric_limits<unsigned>::max() << std::bit_width(unsigned(l ^ r)));
for (; unsigned(l) != shared; l &= l - 1) result -= data[l - 1];
for (; unsigned(r) != shared; r &= r - 1) result += data[r - 1];
return result;
}
T prod(int l, int r) const { return prod_as<T>(l, r); }
};
#line 7 "data_structure/weighted_wavelet_matrix.hpp"
template <class T = std::int64_t, class Count = std::int64_t,
class Sum = std::conditional_t<(sizeof(T) < 8), std::int64_t, __int128_t>>
class WeightedWaveletMatrix {
static_assert(std::is_integral_v<T> && !std::is_same_v<T, bool> && sizeof(T) <= 8);
static_assert(std::is_integral_v<Count> && std::is_signed_v<Count> && sizeof(Count) <= 8);
static_assert(std::numeric_limits<Sum>::is_integer && std::numeric_limits<Sum>::is_signed);
struct Node {
Count count{};
Sum sum{};
Node &operator+=(const Node &x) { count += x.count; sum += x.sum; return *this; }
Node &operator-=(const Node &x) { count -= x.count; sum -= x.sum; return *this; }
friend Node operator-(Node a, const Node &b) { return a -= b; }
};
wavelet_detail::ValueIndex<T> index;
std::vector<Count> weights;
std::vector<QuadBitVector> bits;
std::vector<std::array<int, 4>> middle;
std::vector<FenwickSum<Count>> counts;
std::vector<FenwickSum<Sum>> sums;
bool sum_ready = false;
int levels = 0;
void check([[maybe_unused]] int l, [[maybe_unused]] int r) const { assert(0 <= l && l <= r && r <= size()); }
template <bool WithSum>
Node row(int d, int l, int r) const {
if constexpr (WithSum) return {counts[d].prod(l, r), sums[d].prod(l, r)};
else return {counts[d].prod(l, r), Sum(0)};
}
void add(int d, int p, Node delta) {
counts[d].add(p, delta.count);
if (sum_ready) sums[d].add(p, delta.sum);
}
void add_pair(int d, int p, Node a, int q, Node b) {
counts[d].add_pair(p, a.count, q, b.count);
if (sum_ready) sums[d].add_pair(p, a.sum, q, b.sum);
}
void prepare_sum_impl() {
if (sum_ready) return;
sums.resize(counts.size());
std::vector<Sum> current(size()), next(size());
for (int i = 0; i < size(); ++i) current[i] = Sum(index.values[index.original[i]]) * weights[i];
sums[0].build(current);
for (int d = 0; d < levels; ++d) {
std::array<std::vector<Sum>, 3> zero;
zero[0].reserve(middle[d][2]);
zero[1].reserve(middle[d][1]);
zero[2].reserve(middle[d][3] - middle[d][2]);
auto cursor = middle[d];
for (int i = 0; i < size(); ++i) {
const auto &block = bits[d].blocks[i >> 6];
int digit = ((block.low >> (i & 63)) & 1) | (((block.high >> (i & 63)) & 1) << 1);
if (!(digit & 2)) zero[0].push_back(current[i]);
if (!(digit & 1)) zero[1 + (digit >> 1)].push_back(current[i]);
next[cursor[digit]++] = current[i];
}
for (int j = 0; j < 3; ++j) sums[3 * d + 1 + j].build(zero[j]);
current.swap(next);
}
sum_ready = true;
}
template <bool Single = false, bool WithSum = true, class Choose>
std::pair<std::uint32_t, Node> walk(int l, int r, Node total, Choose choose) const {
std::uint32_t x = 0;
for (int d = 0; d < levels && total.count; ++d) {
if constexpr (Single) if (total.count == 1) return {x, total};
int base = 3 * d + 1;
int a = bits[d].below(l, 2), b = bits[d].below(r, 2);
Node zero = row<WithSum>(base, a, b);
bool high = choose(zero, total);
total = high ? total - zero : zero;
int high_l = high ? l - a : a, high_r = high ? r - b : b;
int c = bits[d].rank(l, 2 * high), e = bits[d].rank(r, 2 * high);
zero = row<WithSum>(base + 1 + high, c, e);
bool low = choose(zero, total);
total = low ? total - zero : zero;
int digit = 2 * high + low;
x = (x << 2) | digit;
l = middle[d][digit] + (low ? high_l - c : c);
r = middle[d][digit] + (low ? high_r - e : e);
}
return {x, total};
}
template <bool WithSum = false>
Node less_than(int l, int r, std::uint32_t x, int first = 0) const {
if (!x || l == r) return {};
if constexpr (WithSum) assert(sum_ready);
if (x >= index.values.size()) return row<WithSum>(0, l, r);
Node answer{};
for (int d = first; d < levels && l < r; ++d) {
int digit = (x >> (2 * (levels - 1 - d))) & 3, high = digit >> 1, low = digit & 1;
int a = bits[d].below(l, 2), b = bits[d].below(r, 2), base = 3 * d + 1;
if (high) answer += row<WithSum>(base, a, b);
int c = bits[d].rank(l, 2 * high), e = bits[d].rank(r, 2 * high);
if (low) answer += row<WithSum>(base + 1 + high, c, e);
l = middle[d][digit] + (low ? (high ? l - a : a) - c : c);
r = middle[d][digit] + (low ? (high ? r - b : b) - e : e);
}
return answer;
}
template <bool WithSum = false>
Node between(int l, int r, std::uint32_t low, std::uint32_t high) const {
if (low == high || l == r) return {};
if (!low || high == index.values.size()) return less_than<WithSum>(l, r, high) - less_than<WithSum>(l, r, low);
int common = levels - (std::bit_width(low ^ high) + 1) / 2;
for (int d = 0; d < common; ++d) {
int digit = (low >> (2 * (levels - 1 - d))) & 3;
l = middle[d][digit] + bits[d].rank(l, digit);
r = middle[d][digit] + bits[d].rank(r, digit);
}
return less_than<WithSum>(l, r, high, common) - less_than<WithSum>(l, r, low, common);
}
template <bool Largest>
Sum selected_sum(int l, int r, Count k, Node total) const {
assert(0 <= k && k <= total.count);
if (!k) return 0;
if (k == total.count) return total.sum;
Sum answer = 0;
auto [x, leaf] = walk<true>(l, r, total, [&](Node zero, Node all) {
Node preferred = Largest ? all - zero : zero;
bool skip = k >= preferred.count;
if (skip) k -= preferred.count, answer += preferred.sum;
return Largest != skip;
});
return answer + (leaf.count == 1 ? leaf.sum : Sum(index.values[x])) * k;
}
template <bool Largest, bool Lower>
Count sum_bound(int l, int r, Sum target) const {
check(l, r);
if constexpr (Lower) {
if (target <= 0) return 0;
} else if (target < 0) return -1;
if (l == r) return Lower ? -1 : 0;
assert(index.values.front() >= 0);
assert(sum_ready);
Node total = row<true>(0, l, r);
if constexpr (Lower) {
if (total.sum < target) return -1;
} else if (total.sum <= target) return total.count;
Count answer = 0;
auto [x, leaf] = walk<true>(l, r, total, [&](Node zero, Node all) {
Node preferred = Largest ? all - zero : zero;
bool skip = Lower ? preferred.sum < target : preferred.sum <= target;
if (skip) target -= preferred.sum, answer += preferred.count;
return Largest != skip;
});
Sum value = leaf.count == 1 ? leaf.sum : Sum(index.values[x]);
return answer + Count(target / value) + (Lower && target % value != 0);
}
public:
using value_type = T;
using count_type = Count;
using sum_type = Sum;
WeightedWaveletMatrix() { build({}, {}); }
explicit WeightedWaveletMatrix(const std::vector<T> &a) { build(a, std::vector<Count>(a.size(), 1)); }
WeightedWaveletMatrix(const std::vector<T> &a, const std::vector<Count> &w, bool with_sum = true) { build(a, w, with_sum); }
void prepare_sum() { prepare_sum_impl(); }
void build(const std::vector<T> &a, const std::vector<Count> &w, bool with_sum = true) {
assert(a.size() == w.size() && a.size() <= std::size_t(std::numeric_limits<int>::max()));
index.build(a);
weights = w;
levels = index.values.size() <= 1 ? 0 : (std::bit_width(unsigned(index.values.size() - 1)) + 1) / 2;
bits.clear();
bits.reserve(levels);
middle.resize(levels);
counts.clear();
counts.resize(3 * levels + 1);
sums.clear();
sum_ready = false;
std::vector<Count> root = w;
for (int i = 0; i < size(); ++i) {
assert(w[i] >= 0);
}
counts[0].build(root);
std::vector<std::uint32_t> current = index.original, next(a.size());
std::vector<Count> next_weights(a.size());
for (int d = 0; d < levels; ++d) {
int shift = 2 * (levels - 1 - d);
bits.emplace_back(size());
bits.back().build(current, shift);
auto totals = bits.back().ranks(size());
middle[d] = {0, totals[0], totals[0] + totals[1], totals[0] + totals[1] + totals[2]};
std::array<std::vector<Count>, 3> zero;
zero[0].reserve(middle[d][2]);
zero[1].reserve(totals[0]);
zero[2].reserve(totals[2]);
auto cursor = middle[d];
for (int i = 0; i < size(); ++i) {
int digit = (current[i] >> shift) & 3;
Count node = root[i];
if (!(digit & 2)) zero[0].push_back(node);
if (!(digit & 1)) zero[1 + (digit >> 1)].push_back(node);
int p = cursor[digit]++;
next[p] = current[i], next_weights[p] = node;
}
for (int j = 0; j < 3; ++j) counts[3 * d + 1 + j].build(zero[j]);
current.swap(next), root.swap(next_weights);
}
if (with_sum) prepare_sum();
}
void build(const std::vector<T> &a) { build(a, std::vector<Count>(a.size(), 1)); }
int size() const { return index.n; }
bool empty() const { return size() == 0; }
T access(int p) const {
assert(0 <= p && p < size());
return index.values[index.original[p]];
}
T operator[](int p) const { return access(p); }
Count weight(int p) const {
assert(0 <= p && p < size());
return weights[p];
}
void add_weight(int p, Count delta) {
assert(0 <= p && p < size());
assert(delta >= -weights[p] && (delta <= 0 || weights[p] <= std::numeric_limits<Count>::max() - delta));
if (!delta) return;
weights[p] += delta;
std::uint32_t x = index.original[p];
Node change{delta, sum_ready ? Sum(index.values[x]) * delta : Sum(0)};
add(0, p, change);
for (int d = 0; d < levels; ++d) {
int digit = (x >> (2 * (levels - 1 - d))) & 3;
int rank = bits[d].rank(p, digit);
if (!(digit & 2)) add(3 * d + 1, bits[d].below(p, 2), change);
if (!(digit & 1)) add(3 * d + 2 + (digit >> 1), rank, change);
p = middle[d][digit] + rank;
}
}
void set_weight(int p, Count value) { assert(value >= 0); add_weight(p, value - weight(p)); }
void move_weight(int from, int to, Count amount) {
assert(0 <= from && from < size() && 0 <= to && to < size());
assert(0 <= amount && amount <= weights[from]);
if (from == to || !amount) return;
assert(weights[to] <= std::numeric_limits<Count>::max() - amount);
weights[from] -= amount;
weights[to] += amount;
std::uint32_t x = index.original[from], y = index.original[to];
Node minus{Count(-amount), sum_ready ? -Sum(index.values[x]) * amount : Sum(0)}, plus{amount, sum_ready ? Sum(index.values[y]) * amount : Sum(0)};
add_pair(0, from, minus, to, plus);
for (int d = 0; d < levels; ++d) {
int shift = 2 * (levels - 1 - d), a = (x >> shift) & 3, b = (y >> shift) & 3;
int p = bits[d].rank(from, a), q = bits[d].rank(to, b), base = 3 * d + 1;
if (!(a & 2) && !(b & 2)) add_pair(base, bits[d].below(from, 2), minus, bits[d].below(to, 2), plus);
else if (!(a & 2)) add(base, bits[d].below(from, 2), minus);
else if (!(b & 2)) add(base, bits[d].below(to, 2), plus);
if (!(a & 1) && a == b) add_pair(base + 1 + (a >> 1), p, minus, q, plus);
else {
if (!(a & 1)) add(base + 1 + (a >> 1), p, minus);
if (!(b & 1)) add(base + 1 + (b >> 1), q, plus);
}
from = middle[d][a] + p, to = middle[d][b] + q;
}
}
Count count(int l, int r) const { check(l, r); return counts[0].prod(l, r); }
Count count(int l, int r, const T &x) const {
check(l, r);
auto c = index.code(x);
if (c == index.values.size() || index.values[c] != x) return 0;
return between(l, r, c, c + 1).count;
}
Count range_freq(int l, int r, const T &upper) const { check(l, r); return less_than(l, r, index.code(upper)).count; }
Count range_freq(int l, int r, const T &lower, const T &upper) const {
check(l, r);
return lower < upper ? between(l, r, index.code(lower), index.code(upper)).count : 0;
}
Sum range_sum(int l, int r) const { check(l, r); assert(sum_ready); return sums[0].prod(l, r); }
Sum range_sum(int l, int r, const T &upper) const { check(l, r); return less_than<true>(l, r, index.code(upper)).sum; }
Sum range_sum(int l, int r, const T &lower, const T &upper) const {
check(l, r);
return lower < upper ? between<true>(l, r, index.code(lower), index.code(upper)).sum : 0;
}
T kth_smallest(int l, int r, Count k) const {
check(l, r);
Node total{counts[0].prod(l, r), Sum(0)};
assert(0 <= k && k < total.count);
auto [x, leaf] = walk<false, false>(l, r, total, [&](Node zero, Node) {
bool one = k >= zero.count;
if (one) k -= zero.count;
return one;
});
return index.values[x];
}
T kth_largest(int l, int r, Count k) const { return kth_smallest(l, r, count(l, r) - 1 - k); }
Sum sum_smallest(int l, int r, Count k) const { check(l, r); assert(sum_ready); return selected_sum<false>(l, r, k, row<true>(0, l, r)); }
Sum sum_largest(int l, int r, Count k) const { check(l, r); assert(sum_ready); return selected_sum<true>(l, r, k, row<true>(0, l, r)); }
Count min_count_smallest(int l, int r, Sum target) const { return sum_bound<false, true>(l, r, target); }
Count min_count_largest(int l, int r, Sum target) const { return sum_bound<true, true>(l, r, target); }
Count max_count_smallest(int l, int r, Sum budget) const { return sum_bound<false, false>(l, r, budget); }
Count max_count_largest(int l, int r, Sum budget) const { return sum_bound<true, false>(l, r, budget); }
};
#line 5 "data_structure/point_set_wavelet_matrix.hpp"
template <class T = std::int64_t, class Sum = std::conditional_t<(sizeof(T) < 8), std::int64_t, __int128_t>>
class PointSetWaveletMatrix {
std::vector<T> keys;
std::vector<int> offsets{0}, current;
WeightedWaveletMatrix<T, int, Sum> wm;
void check([[maybe_unused]] int l, [[maybe_unused]] int r) const { assert(0 <= l && l <= r && r <= size()); }
int find(int p, const T &x) const {
assert(0 <= p && p < size());
return std::lower_bound(keys.begin() + offsets[p], keys.begin() + offsets[p + 1], x) - keys.begin();
}
public:
using value_type = T;
using sum_type = Sum;
PointSetWaveletMatrix() = default;
explicit PointSetWaveletMatrix(const std::vector<T> &a, const std::vector<std::pair<int, T>> &candidates = {}) { build(a, candidates); }
void build(const std::vector<T> &a, const std::vector<std::pair<int, T>> &candidates = {}) {
assert(a.size() <= std::size_t(std::numeric_limits<int>::max()));
std::vector<std::pair<int, T>> points = candidates;
points.reserve(a.size() + candidates.size());
for (int i = 0; i < static_cast<int>(a.size()); ++i) points.emplace_back(i, a[i]);
std::sort(points.begin(), points.end());
points.erase(std::unique(points.begin(), points.end()), points.end());
assert(points.size() <= std::size_t(std::numeric_limits<int>::max()));
offsets.assign(a.size() + 1, 0);
keys.clear();
keys.reserve(points.size());
for (const auto &[p, x] : points) {
assert(0 <= p && p < static_cast<int>(a.size()));
++offsets[p + 1];
keys.push_back(x);
}
for (int i = 0; i < static_cast<int>(a.size()); ++i) offsets[i + 1] += offsets[i];
current.resize(a.size());
std::vector<int> weights(points.size());
for (int i = 0; i < size(); ++i) weights[current[i] = find(i, a[i])] = 1;
wm.build(keys, weights);
}
int size() const { return static_cast<int>(current.size()); }
bool empty() const { return current.empty(); }
T access(int p) const { assert(0 <= p && p < size()); return keys[current[p]]; }
T operator[](int p) const { return access(p); }
bool contains(int p, const T &x) const {
int q = find(p, x);
return q < offsets[p + 1] && keys[q] == x;
}
bool set(int p, const T &x) {
int q = find(p, x);
if (q == offsets[p + 1] || keys[q] != x) return false;
wm.move_weight(current[p], q, 1);
current[p] = q;
return true;
}
int count(int l, int r, const T &x) const { check(l, r); return wm.count(offsets[l], offsets[r], x); }
int rank(const T &x, int r) const { return count(0, r, x); }
int range_freq(int l, int r, const T &upper) const { check(l, r); return wm.range_freq(offsets[l], offsets[r], upper); }
int range_freq(int l, int r, const T &lower, const T &upper) const { check(l, r); return wm.range_freq(offsets[l], offsets[r], lower, upper); }
Sum range_sum(int l, int r) const { check(l, r); return wm.range_sum(offsets[l], offsets[r]); }
Sum range_sum(int l, int r, const T &upper) const { check(l, r); return wm.range_sum(offsets[l], offsets[r], upper); }
Sum range_sum(int l, int r, const T &lower, const T &upper) const { check(l, r); return wm.range_sum(offsets[l], offsets[r], lower, upper); }
T kth_smallest(int l, int r, int k) const { check(l, r); return wm.kth_smallest(offsets[l], offsets[r], k); }
T kth_largest(int l, int r, int k) const { check(l, r); return wm.kth_largest(offsets[l], offsets[r], k); }
Sum sum_smallest(int l, int r, int k) const { check(l, r); return wm.sum_smallest(offsets[l], offsets[r], k); }
Sum sum_largest(int l, int r, int k) const { check(l, r); return wm.sum_largest(offsets[l], offsets[r], k); }
int min_count_smallest(int l, int r, Sum target) const { check(l, r); return wm.min_count_smallest(offsets[l], offsets[r], target); }
int min_count_largest(int l, int r, Sum target) const { check(l, r); return wm.min_count_largest(offsets[l], offsets[r], target); }
int max_count_smallest(int l, int r, Sum budget) const { check(l, r); return wm.max_count_smallest(offsets[l], offsets[r], budget); }
int max_count_largest(int l, int r, Sum budget) const { check(l, r); return wm.max_count_largest(offsets[l], offsets[r], budget); }
};