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 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 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 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(0, hi)(rng); } int main() { srand(3); + // Treaps as sets rep(it,0,1000) { vector 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 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 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