proconlib
This documentation is automatically generated by competitive-verifier/competitive-verifier
tests/wavelet_matrix/range_prod.test.cpp
- View this file on GitHub
- Last update: 2025-02-10 20:05:30+09:00
- Problem: https://judge.yosupo.jp/problem/point_add_rectangle_sum
Depends on
- Bit vector (tools/bit_vector.hpp)
- Bit width required to represent a given integer (tools/bit_width.hpp)
- $\left\lfloor \log_2(x) \right\rfloor$ (tools/floor_log2.hpp)
- std::is_integral extended for tools::int128_t and tools::uint128_t (tools/is_integral.hpp)
- std::is_signed extended for tools::int128_t and tools::uint128_t (tools/is_signed.hpp)
- std::less by first (tools/less_by_first.hpp)
- std::lower_bound, but returns index (tools/lower_bound.hpp)
- std::make_unsigned extended for tools::int128_t and tools::uint128_t (tools/make_unsigned.hpp)
- Wavelet matrix (tools/wavelet_matrix.hpp)
Code
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/point_add_rectangle_sum
#include <iostream>
#include <vector>
#include <queue>
#include <tuple>
#include <iterator>
#include "atcoder/fenwicktree.hpp"
#include "tools/wavelet_matrix.hpp"
using ll = long long;
int main() {
std::cin.tie(nullptr);
std::ios_base::sync_with_stdio(false);
int N, Q;
std::cin >> N >> Q;
tools::wavelet_matrix<int> wm;
std::vector<int> ws;
std::queue<int> qts;
std::queue<std::tuple<int, int, int, int>> q2s;
for (int i = 0; i < N; ++i) {
int x, y, w;
std::cin >> x >> y >> w;
wm.add_point(x, y);
ws.push_back(w);
}
for (int q = 0; q < Q; ++q) {
int qt;
std::cin >> qt;
if (qt == 0) {
int x, y, w;
std::cin >> x >> y >> w;
wm.add_point(x, y);
ws.push_back(w);
qts.push(0);
} else {
int l, d, r, u;
std::cin >> l >> d >> r >> u;
qts.push(1);
q2s.emplace(l, d, r, u);
}
}
std::vector<atcoder::fenwick_tree<ll>> fws;
for (const auto& v : wm.build()) {
fws.emplace_back(v.size());
for (int i = 0; i < std::ssize(v); ++i) {
if (int(v[i]) < N) {
fws.back().add(i, ws[v[i]]);
}
}
}
int p = N;
for (int q = 0; q < Q; ++q) {
const auto qt = qts.front();
qts.pop();
if (qt == 0) {
for (const auto& [h, i] : wm.apply(p)) {
fws[h].add(i, ws[p]);
}
++p;
} else {
const auto [l, d, r, u] = q2s.front();
q2s.pop();
ll answer = 0;
for (const auto& [h, il, ir] : wm.range_prod(l, r, u)) {
answer += fws[h].sum(il, ir);
}
for (const auto& [h, il, ir] : wm.range_prod(l, r, d)) {
answer -= fws[h].sum(il, ir);
}
std::cout << answer << '\n';
}
}
return 0;
}
#line 1 "tests/wavelet_matrix/range_prod.test.cpp"
// competitive-verifier: PROBLEM https://judge.yosupo.jp/problem/point_add_rectangle_sum
#include <iostream>
#include <vector>
#include <queue>
#include <tuple>
#include <iterator>
#line 1 "lib/ac-library/atcoder/fenwicktree.hpp"
#include <cassert>
#line 6 "lib/ac-library/atcoder/fenwicktree.hpp"
#line 1 "lib/ac-library/atcoder/internal_type_traits.hpp"
#line 5 "lib/ac-library/atcoder/internal_type_traits.hpp"
#include <numeric>
#include <type_traits>
namespace atcoder {
namespace internal {
#ifndef _MSC_VER
template <class T>
using is_signed_int128 =
typename std::conditional<std::is_same<T, __int128_t>::value ||
std::is_same<T, __int128>::value,
std::true_type,
std::false_type>::type;
template <class T>
using is_unsigned_int128 =
typename std::conditional<std::is_same<T, __uint128_t>::value ||
std::is_same<T, unsigned __int128>::value,
std::true_type,
std::false_type>::type;
template <class T>
using make_unsigned_int128 =
typename std::conditional<std::is_same<T, __int128_t>::value,
__uint128_t,
unsigned __int128>;
template <class T>
using is_integral = typename std::conditional<std::is_integral<T>::value ||
is_signed_int128<T>::value ||
is_unsigned_int128<T>::value,
std::true_type,
std::false_type>::type;
template <class T>
using is_signed_int = typename std::conditional<(is_integral<T>::value &&
std::is_signed<T>::value) ||
is_signed_int128<T>::value,
std::true_type,
std::false_type>::type;
template <class T>
using is_unsigned_int =
typename std::conditional<(is_integral<T>::value &&
std::is_unsigned<T>::value) ||
is_unsigned_int128<T>::value,
std::true_type,
std::false_type>::type;
template <class T>
using to_unsigned = typename std::conditional<
is_signed_int128<T>::value,
make_unsigned_int128<T>,
typename std::conditional<std::is_signed<T>::value,
std::make_unsigned<T>,
std::common_type<T>>::type>::type;
#else
template <class T> using is_integral = typename std::is_integral<T>;
template <class T>
using is_signed_int =
typename std::conditional<is_integral<T>::value && std::is_signed<T>::value,
std::true_type,
std::false_type>::type;
template <class T>
using is_unsigned_int =
typename std::conditional<is_integral<T>::value &&
std::is_unsigned<T>::value,
std::true_type,
std::false_type>::type;
template <class T>
using to_unsigned = typename std::conditional<is_signed_int<T>::value,
std::make_unsigned<T>,
std::common_type<T>>::type;
#endif
template <class T>
using is_signed_int_t = std::enable_if_t<is_signed_int<T>::value>;
template <class T>
using is_unsigned_int_t = std::enable_if_t<is_unsigned_int<T>::value>;
template <class T> using to_unsigned_t = typename to_unsigned<T>::type;
} // namespace internal
} // namespace atcoder
#line 8 "lib/ac-library/atcoder/fenwicktree.hpp"
namespace atcoder {
// Reference: https://en.wikipedia.org/wiki/Fenwick_tree
template <class T> struct fenwick_tree {
using U = internal::to_unsigned_t<T>;
public:
fenwick_tree() : _n(0) {}
explicit fenwick_tree(int n) : _n(n), data(n) {}
void add(int p, T x) {
assert(0 <= p && p < _n);
p++;
while (p <= _n) {
data[p - 1] += U(x);
p += p & -p;
}
}
T sum(int l, int r) {
assert(0 <= l && l <= r && r <= _n);
return sum(r) - sum(l);
}
private:
int _n;
std::vector<U> data;
U sum(int r) {
U s = 0;
while (r > 0) {
s += data[r - 1];
r -= r & -r;
}
return s;
}
};
} // namespace atcoder
#line 1 "tools/wavelet_matrix.hpp"
#line 5 "tools/wavelet_matrix.hpp"
#include <utility>
#include <cstddef>
#line 8 "tools/wavelet_matrix.hpp"
#include <algorithm>
#line 10 "tools/wavelet_matrix.hpp"
#include <array>
#line 12 "tools/wavelet_matrix.hpp"
#include <optional>
#line 1 "tools/bit_vector.hpp"
#include <cstdint>
#line 6 "tools/bit_vector.hpp"
#include <immintrin.h>
// Source: https://nyaannyaan.github.io/library/data-structure-2d/wavelet-matrix.hpp.html
// License: CC0 1.0 Universal
// Author: Nyaan
namespace tools {
class bit_vector {
private:
using u32 = ::std::uint32_t;
using i64 = ::std::int64_t;
using u64 = ::std::uint64_t;
static constexpr u32 w = 64;
::std::vector<u64> m_block;
::std::vector<u32> m_count;
u32 m_size, m_zeros;
public:
bit_vector() {}
explicit bit_vector(int _n) { init(_n); }
__attribute__((optimize("O3", "unroll-loops"))) void init(int _n) {
m_size = m_zeros = _n;
m_block.resize(m_size / w + 1, 0);
m_count.resize(m_block.size(), 0);
}
u32 size() const { return m_size; }
inline u32 get(u32 i) const { return u32(m_block[i / w] >> (i % w)) & 1u; }
inline void set(u32 i) { m_block[i / w] |= 1LL << (i % w); }
__attribute__((target("popcnt"))) void build() {
for (u32 i = 1; i < m_block.size(); ++i)
m_count[i] = m_count[i - 1] + _mm_popcnt_u64(m_block[i - 1]);
m_zeros = rank0(m_size);
}
u32 zeros() const { return m_zeros; }
inline u32 rank0(u32 i) const { return i - rank1(i); }
__attribute__((target("bmi2,popcnt"))) inline u32 rank1(u32 i) const {
return m_count[i / w] + _mm_popcnt_u64(_bzhi_u64(m_block[i / w], i % w));
}
};
}
#line 1 "tools/lower_bound.hpp"
#line 6 "tools/lower_bound.hpp"
namespace tools {
template <class ForwardIterator, class T>
auto lower_bound(ForwardIterator first, ForwardIterator last, const T& value) {
return ::std::distance(first, ::std::lower_bound(first, last, value));
}
template <class ForwardIterator, class T, class Compare>
auto lower_bound(ForwardIterator first, ForwardIterator last, const T& value, Compare comp) {
return ::std::distance(first, ::std::lower_bound(first, last, value, comp));
}
}
#line 1 "tools/less_by_first.hpp"
#line 5 "tools/less_by_first.hpp"
namespace tools {
class less_by_first {
public:
template <class T1, class T2>
bool operator()(const ::std::pair<T1, T2>& x, const ::std::pair<T1, T2>& y) const {
return x.first < y.first;
}
};
}
#line 1 "tools/floor_log2.hpp"
#line 1 "tools/bit_width.hpp"
#include <bit>
#line 1 "tools/is_integral.hpp"
#line 5 "tools/is_integral.hpp"
namespace tools {
template <typename T>
struct is_integral : ::std::is_integral<T> {};
template <typename T>
inline constexpr bool is_integral_v = ::tools::is_integral<T>::value;
}
#line 1 "tools/is_signed.hpp"
#line 5 "tools/is_signed.hpp"
namespace tools {
template <typename T>
struct is_signed : ::std::is_signed<T> {};
template <typename T>
inline constexpr bool is_signed_v = ::tools::is_signed<T>::value;
}
#line 1 "tools/make_unsigned.hpp"
#line 5 "tools/make_unsigned.hpp"
namespace tools {
template <typename T>
struct make_unsigned : ::std::make_unsigned<T> {};
template <typename T>
using make_unsigned_t = typename ::tools::make_unsigned<T>::type;
}
#line 10 "tools/bit_width.hpp"
namespace tools {
template <typename T>
constexpr int bit_width(T) noexcept;
template <typename T>
constexpr int bit_width(const T x) noexcept {
static_assert(::tools::is_integral_v<T> && !::std::is_same_v<::std::remove_cv_t<T>, bool>);
if constexpr (::tools::is_signed_v<T>) {
assert(x >= 0);
return ::tools::bit_width<::tools::make_unsigned_t<T>>(x);
} else {
return ::std::bit_width(x);
}
}
}
#line 6 "tools/floor_log2.hpp"
namespace tools {
template <typename T>
constexpr T floor_log2(T x) noexcept {
assert(x > 0);
return ::tools::bit_width(x) - 1;
}
}
#line 17 "tools/wavelet_matrix.hpp"
namespace tools {
template <typename T>
class wavelet_matrix {
private:
::std::vector<::std::pair<T, T>> m_ps;
::std::vector<::std::pair<T, ::std::size_t>> m_xs;
::std::vector<T> m_ys;
::std::vector<::std::size_t> m_is;
::std::vector<::tools::bit_vector> m_bvs;
::std::size_t iid(const ::std::size_t i) const {
return ::tools::lower_bound(this->m_xs.begin(), this->m_xs.end(), ::std::make_pair(this->m_ps[i].first, i));
}
::std::size_t xid(const T& x) const {
return ::tools::lower_bound(this->m_xs.begin(), this->m_xs.end(), ::std::make_pair(x, ::std::size_t{}), ::tools::less_by_first());
}
::std::size_t yid(const T& y) const {
return ::tools::lower_bound(this->m_ys.begin(), this->m_ys.end(), y);
}
::std::size_t lg() const {
return this->m_bvs.size();
}
bool built() const {
return this->lg() > 0;
}
public:
wavelet_matrix() = default;
wavelet_matrix(const ::tools::wavelet_matrix<T>&) = default;
wavelet_matrix(::tools::wavelet_matrix<T>&&) = default;
~wavelet_matrix() = default;
::tools::wavelet_matrix<T>& operator=(const ::tools::wavelet_matrix<T>&) = default;
::tools::wavelet_matrix<T>& operator=(::tools::wavelet_matrix<T>&&) = default;
::std::size_t size() const {
return this->m_ps.size();
}
::std::size_t add_point(const T& x, const T& y) {
assert(!this->built());
this->m_ps.emplace_back(x, y);
return this->size() - 1;
}
::std::pair<T, T> get_point(const ::std::size_t i) const {
assert(i < this->size());
return this->m_ps[i];
}
const ::std::vector<::std::pair<T, T>>& points() const {
return this->m_ps;
}
::std::vector<::std::vector<::std::size_t>> build() {
assert(!this->built());
this->m_xs.reserve(this->size());
for (::std::size_t i = 0; i < this->size(); ++i) {
this->m_xs.emplace_back(this->m_ps[i].first, i);
}
::std::sort(this->m_xs.begin(), this->m_xs.end());
this->m_ys.reserve(this->size());
::std::transform(this->m_ps.begin(), this->m_ps.end(), ::std::back_inserter(this->m_ys), [](const auto& p) { return p.second; });
::std::sort(this->m_ys.begin(), this->m_ys.end());
this->m_ys.erase(::std::unique(this->m_ys.begin(), this->m_ys.end()), this->m_ys.end());
const auto n = ::std::max<::std::size_t>(this->size(), 1);
this->m_bvs.assign(::tools::floor_log2(::std::max<::std::size_t>(this->m_ys.size(), 1)) + 1, ::tools::bit_vector(n));
::std::vector<::std::size_t> cur;
cur.reserve(n);
::std::transform(this->m_xs.begin(), this->m_xs.end(), ::std::back_inserter(cur), [&](const auto& p) { return this->yid(this->m_ps[p.second].second); });
cur.resize(n);
::std::vector<::std::size_t> nxt(n);
auto res = ::std::vector(this->lg() + 1, ::std::vector<::std::size_t>(n));
::std::transform(this->m_xs.begin(), this->m_xs.end(), res.back().begin(), [&](const auto& p) { return p.second; });
for (::std::size_t h = this->lg(); h --> 0;) {
for (::std::size_t i = 0; i < n; ++i) {
if ((cur[i] >> h) & 1) {
this->m_bvs[h].set(i);
}
}
this->m_bvs[h].build();
::std::array<::std::size_t, 2> offsets = {0, this->m_bvs[h].zeros()};
for (::std::size_t i = 0; i < n; ++i) {
auto& offset = offsets[this->m_bvs[h].get(i)];
nxt[offset] = cur[i];
res[h][offset] = res[h + 1][i];
++offset;
}
::std::swap(cur, nxt);
}
this->m_is = res.front();
res.pop_back();
return res;
}
::std::vector<::std::pair<::std::size_t, ::std::size_t>> apply(::std::size_t i) const {
assert(this->built());
assert(i < this->size());
::std::vector<::std::pair<::std::size_t, ::std::size_t>> res(this->lg());
i = this->iid(i);
for (::std::size_t h = this->lg(); h --> 0;) {
i = this->m_bvs[h].get(i) ? this->m_bvs[h].zeros() + this->m_bvs[h].rank1(i) : this->m_bvs[h].rank0(i);
res[h] = ::std::make_pair(h, i);
}
return res;
}
::std::size_t kth_smallest(const T& l, const T& r, ::std::size_t k) const {
assert(this->built());
assert(l <= r);
auto lid = this->xid(l);
auto rid = this->xid(r);
assert(k < rid - lid);
for (::std::size_t h = this->lg(); h --> 0;) {
const auto l0 = this->m_bvs[h].rank0(lid);
const auto r0 = this->m_bvs[h].rank0(rid);
if (k < r0 - l0) {
lid = l0;
rid = r0;
} else {
k -= r0 - l0;
lid += this->m_bvs[h].zeros() - l0;
rid += this->m_bvs[h].zeros() - r0;
}
}
return this->m_is[lid + k];
}
::std::size_t kth_largest(const T& l, const T& r, const ::std::size_t k) const {
assert(this->built());
assert(l <= r);
const auto lid = this->xid(l);
const auto rid = this->xid(r);
assert(k < rid - lid);
return this->kth_smallest(l, r, rid - lid - k - 1);
}
::std::vector<::std::tuple<::std::size_t, ::std::size_t, ::std::size_t>> range_prod(const T& l, const T& r, const T& u) const {
assert(this->built());
assert(l <= r);
auto lid = this->xid(l);
auto rid = this->xid(r);
const auto uid = this->yid(u);
::std::vector<::std::tuple<::std::size_t, ::std::size_t, ::std::size_t>> res(this->lg());
for (::std::size_t h = this->lg(); h --> 0;) {
const auto l0 = this->m_bvs[h].rank0(lid);
const auto r0 = this->m_bvs[h].rank0(rid);
if ((uid >> h) & 1) {
res[h] = ::std::make_tuple(h, l0, r0);
lid += this->m_bvs[h].zeros() - l0;
rid += this->m_bvs[h].zeros() - r0;
} else {
res[h] = ::std::make_tuple(h, 0, 0);
lid = l0;
rid = r0;
}
}
return res;
}
::std::size_t range_freq(const T& l, const T& r) const {
assert(this->built());
assert(l <= r);
return this->xid(r) - this->xid(l);
}
::std::size_t range_freq(const T& l, const T& r, const T& u) const {
assert(this->built());
assert(l <= r);
auto lid = this->xid(l);
auto rid = this->xid(r);
const auto uid = this->yid(u);
::std::size_t res = 0;
for (::std::size_t h = this->lg(); h --> 0;) {
const auto l0 = this->m_bvs[h].rank0(lid);
const auto r0 = this->m_bvs[h].rank0(rid);
if ((uid >> h) & 1) {
res += r0 - l0;
lid += this->m_bvs[h].zeros() - l0;
rid += this->m_bvs[h].zeros() - r0;
} else {
lid = l0;
rid = r0;
}
}
return res;
}
::std::size_t range_freq(const T& l, const T& r, const T& d, const T& u) const {
assert(this->built());
assert(l <= r);
assert(d <= u);
return this->range_freq(l, r, u) - this->range_freq(l, r, d);
}
::std::optional<T> prev_value(const T& l, const T& r, const T& u) const {
assert(this->built());
assert(l <= r);
const auto k = this->range_freq(l, r, u);
return k == 0 ? ::std::nullopt : ::std::make_optional(this->m_ps[this->kth_smallest(l, r, k - 1)].second);
}
::std::optional<T> next_value(const T& l, const T& r, const T& d) const {
assert(this->built());
assert(l <= r);
const auto k = this->range_freq(l, r, d);
return k == this->range_freq(l, r) ? ::std::nullopt : ::std::make_optional(this->m_ps[this->kth_smallest(l, r, k)].second);
}
::std::pair<typename ::std::vector<::std::size_t>::const_iterator, typename ::std::vector<::std::size_t>::const_iterator> prev_points(const T& l, const T& r, const T& u) const {
assert(this->built());
assert(l <= r);
auto lid = this->xid(l);
auto rid = this->xid(r);
auto k = this->range_freq(l, r, u);
if (k == 0) return ::std::make_pair(this->m_is.cend(), this->m_is.cend());
--k;
for (::std::size_t h = this->lg(); h --> 0;) {
const auto l0 = this->m_bvs[h].rank0(lid);
const auto r0 = this->m_bvs[h].rank0(rid);
if (k < r0 - l0) {
lid = l0;
rid = r0;
} else {
k -= r0 - l0;
lid += this->m_bvs[h].zeros() - l0;
rid += this->m_bvs[h].zeros() - r0;
}
}
return ::std::make_pair(this->m_is.cbegin() + lid, this->m_is.cbegin() + rid);
}
::std::pair<typename ::std::vector<::std::size_t>::const_iterator, typename ::std::vector<::std::size_t>::const_iterator> next_points(const T& l, const T& r, const T& d) const {
assert(this->built());
assert(l <= r);
auto lid = this->xid(l);
auto rid = this->xid(r);
auto k = this->range_freq(l, r, d);
if (k == rid - lid) return ::std::make_pair(this->m_is.cend(), this->m_is.cend());
for (::std::size_t h = this->lg(); h --> 0;) {
const auto l0 = this->m_bvs[h].rank0(lid);
const auto r0 = this->m_bvs[h].rank0(rid);
if (k < r0 - l0) {
lid = l0;
rid = r0;
} else {
k -= r0 - l0;
lid += this->m_bvs[h].zeros() - l0;
rid += this->m_bvs[h].zeros() - r0;
}
}
return ::std::make_pair(this->m_is.cbegin() + lid, this->m_is.cbegin() + rid);
}
};
}
#line 10 "tests/wavelet_matrix/range_prod.test.cpp"
using ll = long long;
int main() {
std::cin.tie(nullptr);
std::ios_base::sync_with_stdio(false);
int N, Q;
std::cin >> N >> Q;
tools::wavelet_matrix<int> wm;
std::vector<int> ws;
std::queue<int> qts;
std::queue<std::tuple<int, int, int, int>> q2s;
for (int i = 0; i < N; ++i) {
int x, y, w;
std::cin >> x >> y >> w;
wm.add_point(x, y);
ws.push_back(w);
}
for (int q = 0; q < Q; ++q) {
int qt;
std::cin >> qt;
if (qt == 0) {
int x, y, w;
std::cin >> x >> y >> w;
wm.add_point(x, y);
ws.push_back(w);
qts.push(0);
} else {
int l, d, r, u;
std::cin >> l >> d >> r >> u;
qts.push(1);
q2s.emplace(l, d, r, u);
}
}
std::vector<atcoder::fenwick_tree<ll>> fws;
for (const auto& v : wm.build()) {
fws.emplace_back(v.size());
for (int i = 0; i < std::ssize(v); ++i) {
if (int(v[i]) < N) {
fws.back().add(i, ws[v[i]]);
}
}
}
int p = N;
for (int q = 0; q < Q; ++q) {
const auto qt = qts.front();
qts.pop();
if (qt == 0) {
for (const auto& [h, i] : wm.apply(p)) {
fws[h].add(i, ws[p]);
}
++p;
} else {
const auto [l, d, r, u] = q2s.front();
q2s.pop();
ll answer = 0;
for (const auto& [h, il, ir] : wm.range_prod(l, r, u)) {
answer += fws[h].sum(il, ir);
}
for (const auto& [h, il, ir] : wm.range_prod(l, r, d)) {
answer -= fws[h].sum(il, ir);
}
std::cout << answer << '\n';
}
}
return 0;
}
Test cases
| Env | Name | Status | Elapsed | Memory |
|---|---|---|---|---|
| g++ | example_00 |
AC |
4 ms | 4 MB |
| g++ | many_points_00 |
AC |
319 ms | 64 MB |
| g++ | many_queries_00 |
AC |
305 ms | 41 MB |
| g++ | max_random_00 |
AC |
316 ms | 54 MB |
| g++ | max_random_01 |
AC |
322 ms | 53 MB |
| g++ | power_of_2_00 |
AC |
193 ms | 25 MB |
| g++ | power_of_2_01 |
AC |
193 ms | 25 MB |
| g++ | random_00 |
AC |
153 ms | 21 MB |
| g++ | random_01 |
AC |
100 ms | 28 MB |
| g++ | random_02 |
AC |
169 ms | 30 MB |
| g++ | small_00 |
AC |
5 ms | 4 MB |
| g++ | small_01 |
AC |
4 ms | 4 MB |
| g++ | small_02 |
AC |
4 ms | 4 MB |
| g++ | small_03 |
AC |
4 ms | 4 MB |
| g++ | small_04 |
AC |
5 ms | 4 MB |
| g++ | xy_concentrate_00 |
AC |
209 ms | 35 MB |
| g++ | xy_concentrate_01 |
AC |
209 ms | 35 MB |
| g++ | xy_concentrate_02 |
AC |
210 ms | 34 MB |