Update treap.hh

trees/treap.hh

@@ -1,11 +1,28 @@
 // @link https://www.acwing.com/file_system/file/content/whole/index/content/8807719/
 // Everything starts from 1
 
-template <typename T> class treap {
+namespace treap_link {
+    template <typename T> class prev {
+    public:
+        T operator()(const T& a) {
+            return a - 1;
+        }
+    };
+    template <typename T> class next {
+    public:
+        T operator()(const T& a) {
+            return a + 1;
+        }
+    };
+}
+
+template <typename T, typename Compare = less<T>,
+          typename Prev = treap_link::prev<T>,
+          typename Next = treap_link::next<T>>
+class treap {
 
 #define lson fhq[u].l
 #define rson fhq[u].r
-
 private:
     using size_type = size_t;
     using value_type = T;
@@ -17,6 +34,9 @@ private:
         size_type l, r, size;
         value_type val;
     };
+    Compare __compare;
+    Prev __prev;
+    Next __next;
     vector<Node> fhq;
     size_type cnt, root;
     size_type x, y, z;
@@ -39,7 +59,7 @@ private:
     void split(size_type u, value_type val, size_type &x, size_type &y) {
         if (!u) x = y = 0;
         else {
-            if (fhq[u].val <= val) x = u, split(rson, val, rson, y);
+            if (!__compare(val, fhq[u].val)) x = u, split(rson, val, rson, y);
             else y = u, split(lson, val, x, lson);
             pushup(u);
         }
@@ -65,8 +85,15 @@ private:
     }
 
 public:
-    treap() : fhq(1), cnt(0), root(0), _size(0) {}
-    explicit treap(size_type n) : fhq(1), cnt(0), root(0), _size(0) {
+    treap(Compare __compare = std::less<T>(),
+          Prev __prev = treap_link::prev<T>(),
+          Next __next = treap_link::next<T>()) :
+        fhq(1), cnt(0), root(0), _size(0),
+        __compare(__compare), __prev(__prev), __next(__next) {}
+    treap(size_type n, Compare __compare = std::less<T>(),
+          Prev __prev = treap_link::prev<T>(),
+          Next __next = treap_link::next<T>()) :
+        fhq(1), cnt(0), root(0), _size(0) {
         fhq.reserve(n + 1);
     }
 
@@ -90,7 +117,6 @@ public:
     }
 
     void insert(value_type val) {
-        assert(!contains(val));
         ++_size;
         split(root, val, x, y);
         root = merge(merge(x, node(val)), y);
@@ -100,13 +126,13 @@ public:
         assert(contains(val));
         --_size;
         split(root, val, x, z);
-        split(x, val - 1, x, y);
+        split(x, __prev(val), x, y);
         y = merge(fhq[y].l, fhq[y].r);
         root = merge(merge(x, y), z);
     }
 
     size_type index_of(value_type val) {
-        split(root, val - 1, x, y);
+        split(root, __prev(val), x, y);
         size_type res = fhq[x].size + 1;
         root = merge(x, y);
         return res;
@@ -118,7 +144,7 @@ public:
     }
 
     value_type prev_element(value_type val) {
-        split(root, val - 1, x, y);
+        split(root, __prev(val), x, y);
         size_type u = x;
         while (rson) u = rson;
         root = merge(x, y);