diff --git a/content/data-structures/Treap.h b/content/data-structures/Treap.h
index aa1d3da4a..976eefae7 100644
--- a/content/data-structures/Treap.h
+++ b/content/data-structures/Treap.h
@@ -1,67 +1,75 @@
/**
- * Author: someone on Codeforces
- * Date: 2017-03-14
- * Source: folklore
+ * Author: Unknown
+ * Date: Unknown
+ * Source: https://cp-algorithms.com/data_structures/treap.html
* Description: A short self-balancing tree. It acts as a
* sequential container with log-time splits/joins, and
- * is easy to augment with additional data.
- * Time: $O(\log N)$
- * Status: stress-tested
+ * is easy to augment with additional data. Also supports
+ * range updates.
+ * Time: $O(\log N)$ per split/merge
+ * Status: stress-tested, tested on cses Substring Reversals and Cut and Paste
*/
#pragma once
struct Node {
Node *l = 0, *r = 0;
- int val, y, c = 1;
+ int val, y, c = 1, rev = 0;
Node(int val) : val(val), y(rand()) {}
- void recalc();
};
int cnt(Node* n) { return n ? n->c : 0; }
-void Node::recalc() { c = cnt(l) + cnt(r) + 1; }
+void pull(Node* n) { if (n) n->c = cnt(n->l) + cnt(n->r)+1; }
+void push(Node* n) {
+ if (!n) return;
+ if (n->rev) {
+ swap(n->l, n->r);
+ if (n->l) n->l->rev ^= 1;
+ if (n->r) n->r->rev ^= 1;
+ n->rev = 0;
+ }
+}
template<class F> void each(Node* n, F f) {
- if (n) { each(n->l, f); f(n->val); each(n->r, f); }
+ if (n) { push(n);each(n->l, f);f(n->val);each(n->r, f); }
}
-pair<Node*, Node*> split(Node* n, int k) {
- if (!n) return {};
- if (cnt(n->l) >= k) { // "n->val >= k" for lower_bound(k)
- auto pa = split(n->l, k);
- n->l = pa.second;
- n->recalc();
- return {pa.first, n};
- } else {
- auto pa = split(n->r, k - cnt(n->l) - 1); // and just "k"
- n->r = pa.first;
- n->recalc();
- return {n, pa.second};
- }
+// Put i first nodes into l, the rest into r
+void split(Node* x, Node*& l, Node*& r, int i) {
+ if (!x) return void(l = r = 0);
+ push(x);
+ // replace cnt(x->l) with x->val for lower_bound(i)
+ if (i <= cnt(x->l)) split(x->l, l, x->l, i), r = x;
+ // and just i instead
+ else split(x->r, x->r, r, i - cnt(x->l) - 1), l = x;
+ pull(x);
}
-Node* merge(Node* l, Node* r) {
- if (!l) return r;
- if (!r) return l;
- if (l->y > r->y) {
- l->r = merge(l->r, r);
- l->recalc();
- return l;
- } else {
- r->l = merge(l, r->l);
- r->recalc();
- return r;
- }
+// Append r to l, store it in x
+void merge(Node*& x, Node* l, Node* r) {
+ push(l), push(r);
+ if (!l || !r) x = l ? l : r;
+ else if (l->y < r->y) merge(r->l, l, r->l), x = r;
+ else merge(l->r, l->r, r), x = l;
+ pull(x);
}
-Node* ins(Node* t, Node* n, int pos) {
- auto pa = split(t, pos);
- return merge(merge(pa.first, n), pa.second);
+void insert(Node*& t, Node* n, int pos) {
+ Node* l, * r;
+ split(t, l, r, pos), merge(l, l, n), merge(t, l, r);
}
// Example application: move the range [l, r) to index k
void move(Node*& t, int l, int r, int k) {
- Node *a, *b, *c;
- tie(a,b) = split(t, l); tie(b,c) = split(b, r - l);
- if (k <= l) t = merge(ins(a, b, k), c);
- else t = merge(a, ins(c, b, k - r));
+ Node* a, * b, * c;
+ split(t, a, c, r), split(a, a, b, l), merge(t, a, c);
+ if (k<=l) insert(t, b, k);
+ else insert(t, b, k - r + l);
+}
+
+// Reverse the range [l, r)
+void rev(Node*& t, int l, int r) {
+ Node* a, * b, * c;
+ split(t, a, c, r), split(a, a, b, l);
+ b->rev ^= 1;
+ merge(a, a, b), merge(t, a, c);
}
diff --git a/stress-tests/data-structures/Treap.cpp b/stress-tests/data-structures/Treap.cpp
index 1ee241da3..7b949119a 100644
--- a/stress-tests/data-structures/Treap.cpp
+++ b/stress-tests/data-structures/Treap.cpp
@@ -2,30 +2,23 @@
#include "../../content/data-structures/Treap.h"
-pair<Node*, Node*> split2(Node* n, int v) {
- if (!n) return {};
- if (n->val >= v) {
- auto pa = split2(n->l, v);
- n->l = pa.second;
- n->recalc();
- return {pa.first, n};
- } else {
- auto pa = split2(n->r, v);
- n->r = pa.first;
- n->recalc();
- return {n, pa.second};
- }
+// l will have nodes <= i, r rest
+void split2(Node* x, Node*& l, Node*& r, int i) {
+ if (!x) return void(l = r = 0);
+ push(x);
+ if (i <= x->val) split2(x->l, l, x->l, i), r = x;
+ else split2(x->r, x->r, r, i), l = x;
+ pull(x);
}
-int ra() {
- static unsigned x;
- x *= 4176481;
- x += 193861934;
- return x >> 1;
+mt19937 rng(10);
+int ra(int hi) {
+ return uniform_int_distribution<int>(0, hi)(rng);
}
int main() {
srand(3);
+ // Treaps as sets
rep(it,0,1000) {
vector<Node> nodes;
vi exp;
@@ -33,16 +26,18 @@ int main() {
nodes.emplace_back(i*2+2);
exp.emplace_back(i*2+2);
}
- Node* n = 0;
- rep(i,0,10)
- n = merge(n, &nodes[i]);
+ Node* root = 0;
+ rep(i,0,10) merge(root, root, &nodes[i]);
int v = rand() % 25;
- int left = cnt(split2(n, v).first);
+ Node *d1, *d2;
+ split2(root, d1, d2, v);
+ int left = cnt(d1);
int rleft = (int)(lower_bound(all(exp), v) - exp.begin());
assert(left == rleft);
}
+ // move range
rep(it,0,10000) {
vector<Node> nodes;
vi exp;
@@ -50,15 +45,14 @@ int main() {
rep(i,0,10) exp.emplace_back(i);
Node* n = 0;
rep(i,0,10)
- n = merge(n, &nodes[i]);
+ merge(n, n, &nodes[i]);
- int i = ra() % 11, j = ra() % 11;
+ int i = ra(10), j = ra(10);
if (i > j) swap(i, j);
- int k = ra() % 11;
+ int k = ra(10);
if (i < k && k < j) continue;
move(n, i, j, k);
- // cerr << i << ' ' << j << ' ' << k << endl;
int nk = (k >= j ? k - (j - i) : k);
vi iv(exp.begin() + i, exp.begin() + j);
@@ -67,10 +61,35 @@ int main() {
int ind = 0;
each(n, [&](int x) {
- // cerr << x << ' ';
assert(x == exp[ind++]);
});
- // cerr << endl;
+ }
+
+ // reverse range
+ rep(it,0,10000) {
+ vector<Node> nodes;
+ vi exp;
+ rep(i,0,10) nodes.emplace_back(i);
+ rep(i,0,10) exp.emplace_back(i);
+ Node* n = 0;
+ rep(i,0,10)
+ merge(n, n, &nodes[i]);
+
+ int rounds = ra(10);
+ // do multiple rounds to try and break lazy propagation
+ rep(i,0,rounds) {
+ int l = ra(9);
+ int r = ra(9);
+ if (r<l) swap(l,r);
+ if (l==r) r++;
+ reverse(exp.begin()+l, exp.begin()+r);
+ rev(n, l, r);
+ }
+
+ int ind = 0;
+ each(n, [&](int x) {
+ assert(x == exp[ind++]);
+ });
}
cout<<"Tests passed!"<<endl;
}