wavelet_matrix.hpp¶
静态半开区间第 \(k\) 小、精确频次、值域计数/求和及前驱后继；排序下标并复用缓冲构建，权重使用各输入到 Sum 的转换。
Complexity: Time: \(O(n \log (n + 1))\) build, \(O(\log (\sigma + 2))\) query. Space: \(O((n + 1) \log (\sigma + 2))\), with \(O(1)\) value/Sum operations.
AC 记录：range_kth_smallest, static_range_sum_with_upper_bound。
Implementation¶
当前头文件，省略 include guard；依赖见 #include。
/// @complexity Time: O(n log(n+1)) build, O(log(sigma+2)) query.
/// Space: O((n+1) log(sigma+2)), with O(1) value/Sum operations.
#include <algorithm>
#include <bit>
#include <cassert>
#include <limits>
#include <numeric>
#include <optional>
#include <utility>
#include <vector>
namespace noya {
/// @brief Coordinate-compressed static wavelet matrix with range order
/// statistics and value-filtered sums. Values use strict-weak-order equivalence;
/// queries return a representative of each class, while sums use each input's
/// own Sum conversion. Sum is an additive commutative group; arithmetic must fit.
/// The input size fits int. Construction owns its values.
template <class T, class Sum = long long> struct wavelet_matrix {
int n = 0;
int lg = 0;
std::vector<T> xs;
std::vector<int> mid;
std::vector<std::vector<int>> s1;
std::vector<std::vector<Sum>> s0;
std::vector<Sum> s;
wavelet_matrix() = default;
explicit wavelet_matrix(const std::vector<T> &a) { build(a); }
/// @brief Rebuild the static matrix from values.
void build(const std::vector<T> &a) {
assert(a.size() <= std::size_t(std::numeric_limits<int>::max()));
n = int(a.size());
std::vector<T> values;
values.reserve(n);
std::vector<int> rk(n), end{0};
end.reserve(a.size() + 1);
{
std::vector<int> order(n);
std::iota(order.begin(), order.end(), 0);
std::stable_sort(order.begin(), order.end(),
[&](int i, int j) { return a[i] < a[j]; });
for (int i = 0; i < n;) {
int j = i + 1;
while (j < n && !(a[order[i]] < a[order[j]]))
++j;
int rank = int(values.size());
values.push_back(a[order[i]]);
for (int k = i; k < j; ++k)
rk[order[k]] = rank;
end.push_back(j);
i = j;
}
}
int m = int(values.size());
lg = std::max(1, int(std::bit_width(unsigned(std::max(0, m - 1)))));
mid.resize(lg);
s1.resize(lg);
s0.resize(lg);
s.resize(a.size() + 1);
s[0] = Sum{};
std::vector<int> next_rank(n);
std::vector<Sum> weight(n), next_weight(n);
for (int i = 0; i < n; ++i) {
weight[i] = Sum(a[i]);
s[i + 1] = s[i] + weight[i];
}
// Finish reading a before replacing xs, including for build(xs).
xs = std::move(values);
for (int dep = 0; dep < lg; ++dep) {
int bit = lg - 1 - dep, zeros = 0;
std::size_t width = std::size_t(1) << bit;
// Rank-class sizes determine the zero partition before scanning a row.
for (std::size_t i = 0; i < std::size_t(m); i += 2 * width)
zeros += end[std::min(i + width, std::size_t(m))] - end[i];
mid[dep] = zeros;
auto &ones = s1[dep];
auto &sums = s0[dep];
ones.resize(std::size_t(n) + 1);
sums.resize(std::size_t(n) + 1);
ones[0] = 0;
sums[0] = Sum{};
int at[2] = {0, zeros};
for (int i = 0; i < n; ++i) {
int one = (rk[i] >> bit) & 1;
ones[i + 1] = ones[i] + one;
sums[i + 1] = sums[i] + (one ? Sum{} : weight[i]);
int j = at[one]++;
next_rank[j] = rk[i];
next_weight[j] = weight[i];
}
rk.swap(next_rank);
weight.swap(next_weight);
}
}
/// @brief Return the number of stored values.
int size() const { return n; }
/// @brief Return the k-th smallest value in values[l, r), where k is
/// zero-indexed.
T kth_smallest(int l, int r, int k) const {
assert(0 <= l && l <= r && r <= n);
assert(0 <= k && k < r - l);
unsigned rk = 0;
for (int dep = 0; dep < lg; dep++) {
int ol = s1[dep][l];
int or_ = s1[dep][r];
int zl = l - ol;
int zr = r - or_;
int c0 = zr - zl;
if (k < c0) {
l = zl;
r = zr;
} else {
k -= c0;
rk |= 1U << (lg - 1 - dep);
l = mid[dep] + ol;
r = mid[dep] + or_;
}
}
assert(rk < xs.size());
return xs[rk];
}
/// @brief Count values x at positions [l, r) satisfying x < hi.
int count_less(int l, int r, const T &hi) const {
check_range(l, r);
if (l == r)
return 0;
int rk = int(std::lower_bound(xs.begin(), xs.end(), hi) - xs.begin());
return count_less_rank(l, r, rk);
}
/// @brief Count values x at positions [l, r) satisfying lo <= x < hi.
int range_freq(int l, int r, const T &lo, const T &hi) const {
assert(!(hi < lo));
return count_less(l, r, hi) - count_less(l, r, lo);
}
/// @brief Count occurrences of val at positions [l, r).
int count(int l, int r, const T &val) const {
check_range(l, r);
if (l == r)
return 0;
int rank = int(std::lower_bound(xs.begin(), xs.end(), val) - xs.begin());
if (rank == int(xs.size()) || val < xs[rank])
return 0;
for (int dep = 0; dep < lg && l < r; ++dep) {
int ol = s1[dep][l], or_ = s1[dep][r];
if ((rank >> (lg - 1 - dep)) & 1) {
l = mid[dep] + ol;
r = mid[dep] + or_;
} else {
l -= ol;
r -= or_;
}
}
return r - l;
}
/// @brief Sum values x at positions [l, r) satisfying lo <= x < hi.
Sum range_sum(int l, int r, const T &lo, const T &hi) const {
assert(!(hi < lo));
return sum_less(l, r, hi) - sum_less(l, r, lo);
}
/// @brief Return the largest value below hi at positions [l, r), or nullopt.
std::optional<T> prev_value(int l, int r, const T &hi) const {
int cnt = count_less(l, r, hi);
if (cnt == 0) {
return std::nullopt;
}
return kth_smallest(l, r, cnt - 1);
}
/// @brief Return the smallest value at least lo at positions [l, r), or
/// nullopt.
std::optional<T> next_value(int l, int r, const T &lo) const {
int cnt = count_less(l, r, lo);
if (cnt == r - l) {
return std::nullopt;
}
return kth_smallest(l, r, cnt);
}
private:
void check_range(int l, int r) const { assert(0 <= l && l <= r && r <= n); }
int count_less_rank(int l, int r, int rk) const {
check_range(l, r);
if (rk <= 0) {
return 0;
}
if (rk >= int(xs.size())) {
return r - l;
}
int res = 0;
for (int dep = 0; dep < lg; dep++) {
int ol = s1[dep][l];
int or_ = s1[dep][r];
int zl = l - ol;
int zr = r - or_;
int bit = lg - 1 - dep;
if ((rk >> bit) & 1) {
res += zr - zl;
l = mid[dep] + ol;
r = mid[dep] + or_;
} else {
l = zl;
r = zr;
}
}
return res;
}
Sum sum_less(int l, int r, const T &hi) const {
check_range(l, r);
if (l == r)
return Sum{};
int rk = int(std::lower_bound(xs.begin(), xs.end(), hi) - xs.begin());
if (rk <= 0) {
return Sum{};
}
if (rk >= int(xs.size())) {
return s[r] - s[l];
}
Sum res{};
for (int dep = 0; dep < lg; dep++) {
int ol = s1[dep][l];
int or_ = s1[dep][r];
int zl = l - ol;
int zr = r - or_;
int bit = lg - 1 - dep;
if ((rk >> bit) & 1) {
res += s0[dep][r] - s0[dep][l];
l = mid[dep] + ol;
r = mid[dep] + or_;
} else {
l = zl;
r = zr;
}
}
return res;
}
};
} // namespace noya
// NOYA_WAVELET_MATRIX_HPP
#ifndef NOYA_WAVELET_MATRIX_HPP
#define NOYA_WAVELET_MATRIX_HPP 1
/// @complexity Time: O(n log(n+1)) build, O(log(sigma+2)) query.
/// Space: O((n+1) log(sigma+2)), with O(1) value/Sum operations.
#include <algorithm>
#include <bit>
#include <cassert>
#include <limits>
#include <numeric>
#include <optional>
#include <utility>
#include <vector>
namespace noya {
/// @brief Coordinate-compressed static wavelet matrix with range order
/// statistics and value-filtered sums. Values use strict-weak-order equivalence;
/// queries return a representative of each class, while sums use each input's
/// own Sum conversion. Sum is an additive commutative group; arithmetic must fit.
/// The input size fits int. Construction owns its values.
template <class T, class Sum = long long> struct wavelet_matrix {
int n = 0;
int lg = 0;
std::vector<T> xs;
std::vector<int> mid;
std::vector<std::vector<int>> s1;
std::vector<std::vector<Sum>> s0;
std::vector<Sum> s;
wavelet_matrix() = default;
explicit wavelet_matrix(const std::vector<T> &a) { build(a); }
/// @brief Rebuild the static matrix from values.
void build(const std::vector<T> &a) {
assert(a.size() <= std::size_t(std::numeric_limits<int>::max()));
n = int(a.size());
std::vector<T> values;
values.reserve(n);
std::vector<int> rk(n), end{0};
end.reserve(a.size() + 1);
{
std::vector<int> order(n);
std::iota(order.begin(), order.end(), 0);
std::stable_sort(order.begin(), order.end(),
[&](int i, int j) { return a[i] < a[j]; });
for (int i = 0; i < n;) {
int j = i + 1;
while (j < n && !(a[order[i]] < a[order[j]]))
++j;
int rank = int(values.size());
values.push_back(a[order[i]]);
for (int k = i; k < j; ++k)
rk[order[k]] = rank;
end.push_back(j);
i = j;
}
}
int m = int(values.size());
lg = std::max(1, int(std::bit_width(unsigned(std::max(0, m - 1)))));
mid.resize(lg);
s1.resize(lg);
s0.resize(lg);
s.resize(a.size() + 1);
s[0] = Sum{};
std::vector<int> next_rank(n);
std::vector<Sum> weight(n), next_weight(n);
for (int i = 0; i < n; ++i) {
weight[i] = Sum(a[i]);
s[i + 1] = s[i] + weight[i];
}
// Finish reading a before replacing xs, including for build(xs).
xs = std::move(values);
for (int dep = 0; dep < lg; ++dep) {
int bit = lg - 1 - dep, zeros = 0;
std::size_t width = std::size_t(1) << bit;
// Rank-class sizes determine the zero partition before scanning a row.
for (std::size_t i = 0; i < std::size_t(m); i += 2 * width)
zeros += end[std::min(i + width, std::size_t(m))] - end[i];
mid[dep] = zeros;
auto &ones = s1[dep];
auto &sums = s0[dep];
ones.resize(std::size_t(n) + 1);
sums.resize(std::size_t(n) + 1);
ones[0] = 0;
sums[0] = Sum{};
int at[2] = {0, zeros};
for (int i = 0; i < n; ++i) {
int one = (rk[i] >> bit) & 1;
ones[i + 1] = ones[i] + one;
sums[i + 1] = sums[i] + (one ? Sum{} : weight[i]);
int j = at[one]++;
next_rank[j] = rk[i];
next_weight[j] = weight[i];
}
rk.swap(next_rank);
weight.swap(next_weight);
}
}
/// @brief Return the number of stored values.
int size() const { return n; }
/// @brief Return the k-th smallest value in values[l, r), where k is
/// zero-indexed.
T kth_smallest(int l, int r, int k) const {
assert(0 <= l && l <= r && r <= n);
assert(0 <= k && k < r - l);
unsigned rk = 0;
for (int dep = 0; dep < lg; dep++) {
int ol = s1[dep][l];
int or_ = s1[dep][r];
int zl = l - ol;
int zr = r - or_;
int c0 = zr - zl;
if (k < c0) {
l = zl;
r = zr;
} else {
k -= c0;
rk |= 1U << (lg - 1 - dep);
l = mid[dep] + ol;
r = mid[dep] + or_;
}
}
assert(rk < xs.size());
return xs[rk];
}
/// @brief Count values x at positions [l, r) satisfying x < hi.
int count_less(int l, int r, const T &hi) const {
check_range(l, r);
if (l == r)
return 0;
int rk = int(std::lower_bound(xs.begin(), xs.end(), hi) - xs.begin());
return count_less_rank(l, r, rk);
}
/// @brief Count values x at positions [l, r) satisfying lo <= x < hi.
int range_freq(int l, int r, const T &lo, const T &hi) const {
assert(!(hi < lo));
return count_less(l, r, hi) - count_less(l, r, lo);
}
/// @brief Count occurrences of val at positions [l, r).
int count(int l, int r, const T &val) const {
check_range(l, r);
if (l == r)
return 0;
int rank = int(std::lower_bound(xs.begin(), xs.end(), val) - xs.begin());
if (rank == int(xs.size()) || val < xs[rank])
return 0;
for (int dep = 0; dep < lg && l < r; ++dep) {
int ol = s1[dep][l], or_ = s1[dep][r];
if ((rank >> (lg - 1 - dep)) & 1) {
l = mid[dep] + ol;
r = mid[dep] + or_;
} else {
l -= ol;
r -= or_;
}
}
return r - l;
}
/// @brief Sum values x at positions [l, r) satisfying lo <= x < hi.
Sum range_sum(int l, int r, const T &lo, const T &hi) const {
assert(!(hi < lo));
return sum_less(l, r, hi) - sum_less(l, r, lo);
}
/// @brief Return the largest value below hi at positions [l, r), or nullopt.
std::optional<T> prev_value(int l, int r, const T &hi) const {
int cnt = count_less(l, r, hi);
if (cnt == 0) {
return std::nullopt;
}
return kth_smallest(l, r, cnt - 1);
}
/// @brief Return the smallest value at least lo at positions [l, r), or
/// nullopt.
std::optional<T> next_value(int l, int r, const T &lo) const {
int cnt = count_less(l, r, lo);
if (cnt == r - l) {
return std::nullopt;
}
return kth_smallest(l, r, cnt);
}
private:
void check_range(int l, int r) const { assert(0 <= l && l <= r && r <= n); }
int count_less_rank(int l, int r, int rk) const {
check_range(l, r);
if (rk <= 0) {
return 0;
}
if (rk >= int(xs.size())) {
return r - l;
}
int res = 0;
for (int dep = 0; dep < lg; dep++) {
int ol = s1[dep][l];
int or_ = s1[dep][r];
int zl = l - ol;
int zr = r - or_;
int bit = lg - 1 - dep;
if ((rk >> bit) & 1) {
res += zr - zl;
l = mid[dep] + ol;
r = mid[dep] + or_;
} else {
l = zl;
r = zr;
}
}
return res;
}
Sum sum_less(int l, int r, const T &hi) const {
check_range(l, r);
if (l == r)
return Sum{};
int rk = int(std::lower_bound(xs.begin(), xs.end(), hi) - xs.begin());
if (rk <= 0) {
return Sum{};
}
if (rk >= int(xs.size())) {
return s[r] - s[l];
}
Sum res{};
for (int dep = 0; dep < lg; dep++) {
int ol = s1[dep][l];
int or_ = s1[dep][r];
int zl = l - ol;
int zr = r - or_;
int bit = lg - 1 - dep;
if ((rk >> bit) & 1) {
res += s0[dep][r] - s0[dep][l];
l = mid[dep] + ol;
r = mid[dep] + or_;
} else {
l = zl;
r = zr;
}
}
return res;
}
};
} // namespace noya
#endif // NOYA_WAVELET_MATRIX_HPP
/// @complexity Time: O(n log(n+1)) build, O(log(sigma+2)) query.
/// Space: O((n+1) log(sigma+2)), with O(1) value/Sum operations.
#include <algorithm>
#include <bit>
#include <cassert>
#include <limits>
#include <numeric>
#include <optional>
#include <utility>
#include <vector>
namespace noya {
/// @brief Coordinate-compressed static wavelet matrix with range order
/// statistics and value-filtered sums. Values use strict-weak-order equivalence;
/// queries return a representative of each class, while sums use each input's
/// own Sum conversion. Sum is an additive commutative group; arithmetic must fit.
/// The input size fits int. Construction owns its values.
template <class T, class Sum = long long> struct wavelet_matrix {
int n = 0;
int lg = 0;
std::vector<T> xs;
std::vector<int> mid;
std::vector<std::vector<int>> s1;
std::vector<std::vector<Sum>> s0;
std::vector<Sum> s;
wavelet_matrix() = default;
explicit wavelet_matrix(const std::vector<T> &a) { build(a); }
/// @brief Rebuild the static matrix from values.
void build(const std::vector<T> &a) {
assert(a.size() <= std::size_t(std::numeric_limits<int>::max()));
n = int(a.size());
std::vector<T> values;
values.reserve(n);
std::vector<int> rk(n), end{0};
end.reserve(a.size() + 1);
{
std::vector<int> order(n);
std::iota(order.begin(), order.end(), 0);
std::stable_sort(order.begin(), order.end(),
[&](int i, int j) { return a[i] < a[j]; });
for (int i = 0; i < n;) {
int j = i + 1;
while (j < n && !(a[order[i]] < a[order[j]]))
++j;
int rank = int(values.size());
values.push_back(a[order[i]]);
for (int k = i; k < j; ++k)
rk[order[k]] = rank;
end.push_back(j);
i = j;
}
}
int m = int(values.size());
lg = std::max(1, int(std::bit_width(unsigned(std::max(0, m - 1)))));
mid.resize(lg);
s1.resize(lg);
s0.resize(lg);
s.resize(a.size() + 1);
s[0] = Sum{};
std::vector<int> next_rank(n);
std::vector<Sum> weight(n), next_weight(n);
for (int i = 0; i < n; ++i) {
weight[i] = Sum(a[i]);
s[i + 1] = s[i] + weight[i];
}
// Finish reading a before replacing xs, including for build(xs).
xs = std::move(values);
for (int dep = 0; dep < lg; ++dep) {
int bit = lg - 1 - dep, zeros = 0;
std::size_t width = std::size_t(1) << bit;
// Rank-class sizes determine the zero partition before scanning a row.
for (std::size_t i = 0; i < std::size_t(m); i += 2 * width)
zeros += end[std::min(i + width, std::size_t(m))] - end[i];
mid[dep] = zeros;
auto &ones = s1[dep];
auto &sums = s0[dep];
ones.resize(std::size_t(n) + 1);
sums.resize(std::size_t(n) + 1);
ones[0] = 0;
sums[0] = Sum{};
int at[2] = {0, zeros};
for (int i = 0; i < n; ++i) {
int one = (rk[i] >> bit) & 1;
ones[i + 1] = ones[i] + one;
sums[i + 1] = sums[i] + (one ? Sum{} : weight[i]);
int j = at[one]++;
next_rank[j] = rk[i];
next_weight[j] = weight[i];
}
rk.swap(next_rank);
weight.swap(next_weight);
}
}
/// @brief Return the number of stored values.
int size() const { return n; }
/// @brief Return the k-th smallest value in values[l, r), where k is
/// zero-indexed.
T kth_smallest(int l, int r, int k) const {
assert(0 <= l && l <= r && r <= n);
assert(0 <= k && k < r - l);
unsigned rk = 0;
for (int dep = 0; dep < lg; dep++) {
int ol = s1[dep][l];
int or_ = s1[dep][r];
int zl = l - ol;
int zr = r - or_;
int c0 = zr - zl;
if (k < c0) {
l = zl;
r = zr;
} else {
k -= c0;
rk |= 1U << (lg - 1 - dep);
l = mid[dep] + ol;
r = mid[dep] + or_;
}
}
assert(rk < xs.size());
return xs[rk];
}
/// @brief Count values x at positions [l, r) satisfying x < hi.
int count_less(int l, int r, const T &hi) const {
check_range(l, r);
if (l == r)
return 0;
int rk = int(std::lower_bound(xs.begin(), xs.end(), hi) - xs.begin());
return count_less_rank(l, r, rk);
}
/// @brief Count values x at positions [l, r) satisfying lo <= x < hi.
int range_freq(int l, int r, const T &lo, const T &hi) const {
assert(!(hi < lo));
return count_less(l, r, hi) - count_less(l, r, lo);
}
/// @brief Count occurrences of val at positions [l, r).
int count(int l, int r, const T &val) const {
check_range(l, r);
if (l == r)
return 0;
int rank = int(std::lower_bound(xs.begin(), xs.end(), val) - xs.begin());
if (rank == int(xs.size()) || val < xs[rank])
return 0;
for (int dep = 0; dep < lg && l < r; ++dep) {
int ol = s1[dep][l], or_ = s1[dep][r];
if ((rank >> (lg - 1 - dep)) & 1) {
l = mid[dep] + ol;
r = mid[dep] + or_;
} else {
l -= ol;
r -= or_;
}
}
return r - l;
}
/// @brief Sum values x at positions [l, r) satisfying lo <= x < hi.
Sum range_sum(int l, int r, const T &lo, const T &hi) const {
assert(!(hi < lo));
return sum_less(l, r, hi) - sum_less(l, r, lo);
}
/// @brief Return the largest value below hi at positions [l, r), or nullopt.
std::optional<T> prev_value(int l, int r, const T &hi) const {
int cnt = count_less(l, r, hi);
if (cnt == 0) {
return std::nullopt;
}
return kth_smallest(l, r, cnt - 1);
}
/// @brief Return the smallest value at least lo at positions [l, r), or
/// nullopt.
std::optional<T> next_value(int l, int r, const T &lo) const {
int cnt = count_less(l, r, lo);
if (cnt == r - l) {
return std::nullopt;
}
return kth_smallest(l, r, cnt);
}
private:
void check_range(int l, int r) const { assert(0 <= l && l <= r && r <= n); }
int count_less_rank(int l, int r, int rk) const {
check_range(l, r);
if (rk <= 0) {
return 0;
}
if (rk >= int(xs.size())) {
return r - l;
}
int res = 0;
for (int dep = 0; dep < lg; dep++) {
int ol = s1[dep][l];
int or_ = s1[dep][r];
int zl = l - ol;
int zr = r - or_;
int bit = lg - 1 - dep;
if ((rk >> bit) & 1) {
res += zr - zl;
l = mid[dep] + ol;
r = mid[dep] + or_;
} else {
l = zl;
r = zr;
}
}
return res;
}
Sum sum_less(int l, int r, const T &hi) const {
check_range(l, r);
if (l == r)
return Sum{};
int rk = int(std::lower_bound(xs.begin(), xs.end(), hi) - xs.begin());
if (rk <= 0) {
return Sum{};
}
if (rk >= int(xs.size())) {
return s[r] - s[l];
}
Sum res{};
for (int dep = 0; dep < lg; dep++) {
int ol = s1[dep][l];
int or_ = s1[dep][r];
int zl = l - ol;
int zr = r - or_;
int bit = lg - 1 - dep;
if ((rk >> bit) & 1) {
res += s0[dep][r] - s0[dep][l];
l = mid[dep] + ol;
r = mid[dep] + or_;
} else {
l = zl;
r = zr;
}
}
return res;
}
};
} // namespace noya
// NOYA_WAVELET_MATRIX_HPP