Advertisement
nasarouf

TREAP

Apr 20th, 2017
679
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C++ 2.53 KB | None | 0 0
  1. #include <iostream>
  2. #include <cstdlib>
  3. #include <functional>
  4. using namespace std;
  5.  
  6.  
  7. /* TREAP
  8. based on ubcarchive treap code by luccasiau@gmail.com
  9. */
  10.  
  11. template<class T, class P, class C>
  12. class Treap {
  13. public:
  14.     Treap *l, *r;
  15.     T v;
  16.     P p;
  17.     C lz, count;
  18.     Treap(T v_ = T(), P p_ = P()) : l(NULL), r(NULL), v(v_), p(p_), lz(C()), count(1) {}
  19.     void lazy() {
  20.         if (this == NULL) return;
  21.         /*insert your lazy-propagation here (only use it in implicit)*/
  22.     }
  23.     C getCount() const { return (this == NULL) ? 0 : count; }
  24.     void update() {
  25.         if (this == NULL) return;
  26.         count = l->getCount() + r->getCount() + 1;
  27.         lazy();
  28.         /*update your update*/
  29.     }
  30.     static Treap* merge(Treap *a, Treap *b) { // regular & implicit
  31.         Treap* ret = NULL;
  32.         a->lazy(); b->lazy();
  33.         if (a == NULL) return b;
  34.         else if (b == NULL) return a;
  35.         else if (a->p < b->p) {
  36.             ret = b;  b->l = Treap::merge(a, b->l);
  37.         }
  38.         else {
  39.             ret = a;  a->r = Treap::merge(a->r, b);
  40.         }
  41.         ret->update();
  42.         return ret;
  43.     }
  44.     Treap* operator[](int k) { // 0-indexed. Kth element from left to right
  45.         lazy();
  46.         int lc = l->getCount();
  47.         if (k - lc == 0) return this;
  48.         else if (k < lc) return l->operator[](k);
  49.         else return r->operator[](k - lc - 1);
  50.     }
  51.     Treap* split(Treap *&a, Treap *&b, T key, function<bool(T, T)> toleft = less_equal<T>()) {
  52.         lazy();
  53.         if (!this) a = b = NULL;
  54.         else if (toleft(v, key)) { a = this; r->split(r, b, key, toleft); }
  55.         else { b = this; l->split(a, l, key, toleft); }
  56.         update();
  57.         return a;
  58.     }
  59.     Treap* insert(Treap *x) {
  60.         lazy();
  61.         Treap *l = NULL, *r = NULL;
  62.         split(l, r, x->v);
  63.         return Treap::merge(Treap::merge(l, x), r);
  64.     }
  65.     Treap* remove(T v) {
  66.         lazy();
  67.         Treap *a = NULL, *b = NULL, *aa = NULL, *ab = NULL;
  68.         split(a, b, v);
  69.         a->split(aa, ab, v, less<T>());
  70.         return Treap::merge(aa, b);
  71.     }
  72.     Treap* find(T x, function<bool(T, T)> lessop = less<T>()) {
  73.         if (!lessop(v,x) && !lessop(x,v)) return this;
  74.         else if (lessop(v,x)) return r ? r->find(x, lessop) : NULL;
  75.         else return l ? l->find(x, lessop) : NULL;
  76.     }
  77. };
  78. template<class T, class P, class C>
  79. ostream& operator<<(ostream& os, const Treap<T,P,C>& x){
  80.     if (&x == NULL) return os<<"()";
  81.     return os << "(" << *(x.l) << x.v << *(x.r) << ")";
  82. }
  83.  
  84.  
  85. //example usage:
  86. int main() {
  87.     Treap<int,int,int>* x=NULL;
  88.     for (int i = 10; i >=0; i--) {
  89.         int v = rand();
  90.         x = x->insert(new Treap<int, int, int>(v,rand()));
  91.         cout << (*x) << endl << endl;
  92.         if (i % 3 == 0) {
  93.             x = x->remove(v);
  94.             cout << "after removal" << (*x) << endl << endl;
  95.         }
  96.     }
  97.     return 0;
  98. }
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement