I was recently up-solving AtCoder Beginner Contest 339 and came across a Persistent Segment Tree Problem G — Smaller Sum.
While implementing my own Persistent Segment Tree. I tried to make it as generic as possible.
Code:
#include <bits/stdc++.h>
using namespace std;
#ifdef LOCAL
#include "debug.hpp"
#else
#define dbg(...)
#endif
/**
* Source:
* https://github.com/kth-competitive-programming/kactl/blob/main/content/various/BumpAllocator.h
* Description: When you need to dynamically allocate many objects and don't
* care about freeing them. "new X" otherwise has an overhead of something like
* 0.05us + 16 bytes per allocation. Status: tested
*/
const size_t SZ = (450 << 20); // 450 mb
static char buf[SZ];
class Alloc {
private:
size_t ptr;
public:
Alloc() : ptr(sizeof(buf)) {}
void *alloc(size_t s) {
assert(s < ptr);
return (void *)&buf[ptr -= s];
}
void reset() { ptr = sizeof(buf); }
};
template <typename Info> class Node {
public:
Info info;
Node *left, *right;
};
template <typename Info> class PSegTree {
public:
typedef Node<Info> node_t;
node_t *root;
vector<node_t *> time;
int n;
Alloc ar;
PSegTree(size_t sz) : PSegTree(vector<Info>(sz, Info())) {}
PSegTree(const vector<Info> &info) {
root = nullptr;
ar = Alloc();
n = (int)info.size();
function<node_t *(int, int, int)> build = [&](int p, int l,
int r) -> node_t * {
if (l == r) {
node_t *node = (node_t *)ar.alloc(sizeof(node_t));
node->info = info[l];
node->left = node->right = nullptr;
return node;
}
int m = l + (r - l) / 2;
return pull(build(2 * p + 1, l, m), build(2 * p + 2, m + 1, r));
};
root = build(0, 0, n - 1);
time.push_back(root);
}
node_t *pull(node_t *left, node_t *right) {
node_t *node = (node_t *)ar.alloc(sizeof(node_t));
node->info = (left->info + right->info);
node->left = left, node->right = right;
return node;
}
void modify(int p, const Info &v) {
function<node_t *(node_t *, int, int)> _modify = [&](node_t *c, int l,
int r) -> node_t * {
if (l == r) {
node_t *node = (node_t *)ar.alloc(sizeof(node_t));
node->info = v;
node->left = node->right = nullptr;
return node;
}
int m = l + (r - l) / 2;
node_t *left = c->left, *right = c->right;
if (p <= m) {
left = _modify(left, l, m);
} else {
right = _modify(right, m + 1, r);
}
return pull(left, right);
};
root = _modify(root, 0, n - 1);
time.push_back(root);
}
Info rangeQuery(int t, int x, int y) {
function<Info(node_t *, int, int)> query = [&](node_t *c, int l,
int r) -> Info {
if (y < l or r < x or c == nullptr) {
return Info();
}
if (x <= l and r <= y) {
return c->info;
}
int m = l + (r - l) / 2;
return query(c->left, l, m) + query(c->right, m + 1, r);
};
return query(time[t], 0, n - 1);
}
};
class Sum {
public:
int64_t x = 0;
Sum() : x(0) {}
Sum(int64_t _x) : x(_x) {}
};
Sum operator+(const Sum &lf, const Sum &rt) {
// dbg(lf.x, rt.x, lf.x + rt.x);
return Sum(lf.x + rt.x);
}
void solve() {
int N;
cin >> N;
vector<int> A(N);
for (int i = 0; i < N; i++)
cin >> A[i];
map<int, int, greater<int>> id;
vector<int> a = A;
sort(a.begin(), a.end());
for (auto &e : a) {
if (id.find(e) == id.end()) {
int sz = (int)id.size();
id[e] = sz;
}
}
PSegTree<Sum> seg(id.size());
for (int i = 0; i < N; i++) {
int idx = id[A[i]];
int64_t val = seg.rangeQuery(i, idx, idx).x;
seg.modify(idx, Sum(val + A[i]));
}
int Q;
cin >> Q;
int64_t b = 0;
for (int _ = 0; _ < Q; _++) {
int64_t l, r, x;
cin >> l >> r >> x;
l = (l ^ b), r = (r ^ b), x = (x ^ b);
b = 0;
if (x != 0) {
auto up = id.lower_bound(x);
if (up != id.end()) {
int idx = up->second;
int64_t rt = seg.rangeQuery(r, 0, idx).x;
int64_t lf = seg.rangeQuery(l - 1, 0, idx).x;
b = rt - lf;
}
}
cout << b << '\n';
}
}
int main() {
ios::sync_with_stdio(false);
cin.tie(nullptr);
solve();
return 0;
}