point_add_rec_sum.hpp¶
#include "noya/point_add_rec_sum.hpp"
#ifndef NOYA_POINT_ADD_REC_SUM_HPP
#define NOYA_POINT_ADD_REC_SUM_HPP 1
#include "noya/persistent_segtree.hpp"
namespace noya {
/// @brief Dynamic point-add rectangle-sum with semi-offline rebuilds.
template <class T> struct dynamic_point_add_rectangle_sum {
persistent_segtree<T> seg;
std::vector<std::array<T, 3>> weighted_points;
std::vector<std::array<T, 3>> buckets;
dynamic_point_add_rectangle_sum() {}
dynamic_point_add_rectangle_sum(std::vector<std::array<T, 3>> &points) {
buckets = points;
build();
};
const int B = 6000;
/// @brief Add a weighted point at (x, y) with weight w.
void add_points(T x, T y, T w) {
buckets.push_back({x, y, w});
if (int(buckets.size()) >= B) {
build();
}
}
void build() {
if (!buckets.empty()) {
for (auto &[x, y, w] : buckets) {
weighted_points.push_back({x, y, w});
}
buckets.clear();
}
seg.build(weighted_points);
}
/// @brief Query the sum of weights in rectangle [l, r) x [d, u).
T query(T l, T r, T d, T u) {
T ans = 0;
ans = seg.prod(l, r, d, u);
for (auto &[x, y, w] : buckets) {
if (l <= x && x < r && d <= y && y < u) {
ans += w;
}
}
return ans;
}
};
} // namespace noya
#endif // NOYA_POINT_ADD_REC_SUM_HPP
#include <algorithm>
#include <vector>
namespace noya {
/// @brief Persistent segment tree supporting 2D rectangle sum and k-th queries.
template <class T> struct persistent_segtree {
std::vector<int> rt;
std::vector<int> ls, rs;
int N, M;
std::vector<T> Xs, Ys;
std::vector<T> sum;
int cnt, C;
persistent_segtree() {}
persistent_segtree(std::vector<std::array<T, 3>> &points) {
if (!points.empty())
build(points);
}
void update(int pre, int &p, int l, int r, int x, T v) {
p = cnt++;
ls[p] = ls[pre];
rs[p] = rs[pre];
sum[p] = sum[pre] + v;
if (l + 1 == r)
return;
int m = (l + r) / 2;
if (x < m)
update(ls[pre], ls[p], l, m, x, v);
else
update(rs[pre], rs[p], m, r, x, v);
}
int query_kth(int L, int R, int l, int r, T k) {
if (l + 1 == r) {
return l;
}
int m = (l + r) / 2;
T val = sum[ls[R]] - sum[ls[L]];
if (val >= k) {
return query_kth(ls[L], ls[R], l, m, k);
} else {
return query_kth(rs[L], rs[R], m, r, k - val);
}
}
T query_prefix(int p, int l, int r, int x) {
if (!p)
return 0;
if (l + 1 == r)
return sum[p];
int m = (l + r) / 2;
if (x < m)
return query_prefix(ls[p], l, m, x);
else
return sum[ls[p]] + query_prefix(rs[p], m, r, x);
}
T query_single(int p, int l, int r, int x) {
if (!p)
return 0;
if (l + 1 == r)
return sum[p];
int m = (l + r) / 2;
if (x < m)
return query_single(ls[p], l, m, x);
else
return query_single(rs[p], m, r, x);
}
T query(int L, int R, int ql, int qr, int l, int r) {
if (ql <= l && r <= qr)
return sum[R] - sum[L];
int m = (l + r) / 2;
T res = 0;
if (ql < m)
res += query(ls[L], ls[R], ql, qr, l, m);
if (qr > m)
res += query(rs[L], rs[R], ql, qr, m, r);
return res;
}
/// @brief Query the prefix sum up to coordinate (r, u) in the original space.
T prod(T r, T u) {
r = std::lower_bound(Xs.begin(), Xs.end(), r) - Xs.begin();
u = std::lower_bound(Ys.begin(), Ys.end(), u) - Ys.begin();
assert(0 <= r && r <= N);
assert(0 <= u && u <= M);
if (u > 0)
return query_prefix(rt[r], 0, M, u - 1);
else
return 0;
}
/// @brief Query the sum of weights in rectangle [l, r) x [d, u).
T prod(T l, T r, T d, T u) {
l = std::lower_bound(Xs.begin(), Xs.end(), l) - Xs.begin();
r = std::lower_bound(Xs.begin(), Xs.end(), r) - Xs.begin();
d = std::lower_bound(Ys.begin(), Ys.end(), d) - Ys.begin();
u = std::lower_bound(Ys.begin(), Ys.end(), u) - Ys.begin();
assert(0 <= l && l <= r && r <= N);
assert(0 <= d && d <= u && u <= M);
if (l == r)
return 0;
else if (d == u)
return 0;
else
return query(rt[l], rt[r], d, u, 0, M);
}
/// @brief Return the k-th smallest Y-coordinate among points with X in [l, r).
T kth(T l, T r, T k) {
l = std::lower_bound(Xs.begin(), Xs.end(), l) - Xs.begin();
r = std::lower_bound(Xs.begin(), Xs.end(), r) - Xs.begin();
assert(0 <= l && l <= r && r <= N);
if (sum[rt[r]] - sum[rt[l]] < k) {
return std::numeric_limits<T>::max();
}
int res = query_kth(rt[l], rt[r], 0, M, k);
return Ys[res];
}
/// @brief Build the persistent segment tree from weighted 2D points {x, y, w}.
void build(std::vector<std::array<T, 3>> points) {
C = int(points.size()) * 30;
cnt = 1;
ls.assign(C, 0);
rs.assign(C, 0);
sum.assign(C, 0);
Xs.clear();
Ys.clear();
for (auto &[x, y, w] : points) {
Xs.push_back(x);
Ys.push_back(y);
}
std::sort(Xs.begin(), Xs.end());
std::sort(Ys.begin(), Ys.end());
Xs.erase(std::unique(Xs.begin(), Xs.end()), Xs.end());
Ys.erase(std::unique(Ys.begin(), Ys.end()), Ys.end());
N = int(Xs.size());
std::vector<std::vector<std::array<T, 2>>> add(N);
for (auto &[x, y, w] : points) {
x = std::lower_bound(Xs.begin(), Xs.end(), x) - Xs.begin();
y = std::lower_bound(Ys.begin(), Ys.end(), y) - Ys.begin();
add[x].push_back({y, w});
}
M = int(Ys.size());
rt.assign(N + 1, 0);
rt[0] = cnt++;
for (int i = 0; i < N; i++) {
rt[i + 1] = rt[i];
for (auto &[y, w] : add[i]) {
int new_rt = 0;
update(rt[i + 1], new_rt, 0, M, y, w);
rt[i + 1] = new_rt;
}
}
std::vector<std::vector<std::array<T, 2>>>().swap(add);
}
/// @brief @return vector of (ls, rs, len, value) for all nodes.
std::vector<std::array<int, 4>> all_nodes() {
std::vector<std::array<int, 4>> nodes(cnt, std::array<int, 4>{});
std::vector<bool> vis(cnt);
vis[0] = true;
for (int i = 0; i <= N; i++) {
auto dfs = [&](auto &dfs, int p, int l, int r) -> void {
if (vis[p])
return;
vis[p] = true;
nodes[p] = {ls[p], rs[p], r - l, sum[p]};
int m = (l + r) / 2;
dfs(dfs, ls[p], l, m);
dfs(dfs, rs[p], m, r);
};
dfs(dfs, rt[i], 0, M);
}
return nodes;
}
};
} // namespace noya
namespace noya {
/// @brief Dynamic point-add rectangle-sum with semi-offline rebuilds.
template <class T> struct dynamic_point_add_rectangle_sum {
persistent_segtree<T> seg;
std::vector<std::array<T, 3>> weighted_points;
std::vector<std::array<T, 3>> buckets;
dynamic_point_add_rectangle_sum() {}
dynamic_point_add_rectangle_sum(std::vector<std::array<T, 3>> &points) {
buckets = points;
build();
};
const int B = 6000;
/// @brief Add a weighted point at (x, y) with weight w.
void add_points(T x, T y, T w) {
buckets.push_back({x, y, w});
if (int(buckets.size()) >= B) {
build();
}
}
void build() {
if (!buckets.empty()) {
for (auto &[x, y, w] : buckets) {
weighted_points.push_back({x, y, w});
}
buckets.clear();
}
seg.build(weighted_points);
}
/// @brief Query the sum of weights in rectangle [l, r) x [d, u).
T query(T l, T r, T d, T u) {
T ans = 0;
ans = seg.prod(l, r, d, u);
for (auto &[x, y, w] : buckets) {
if (l <= x && x < r && d <= y && y < u) {
ans += w;
}
}
return ans;
}
};
} // namespace noya