Quellcode durchsuchen

feat(rbtree): upper bound and lower bound

greatbridf vor 1 Jahr
Ursprung
Commit
5287b2b529
3 geänderte Dateien mit 72 neuen und 17 gelöschten Zeilen
  1. 44 17
      gblibstdc++/include/bits/rbtree
  2. 14 0
      gblibstdc++/include/map
  3. 14 0
      gblibstdc++/include/set

+ 44 - 17
gblibstdc++/include/bits/rbtree

@@ -1,6 +1,7 @@
 #ifndef __GBLIBCPP_BITS_RBTREE__
 #define __GBLIBCPP_BITS_RBTREE__
 
+#include <cstddef>
 #include <functional>
 #include <utility>
 #include <memory>
@@ -269,6 +270,15 @@ private:
         node_alloc_traits::deallocate(alloc, nd, 1);
     }
 
+    constexpr void do_insertion(node* parent,
+        node*& child_inserted, node* nd)
+    {
+        nd->parent = parent;
+        child_inserted = nd;
+        this->balance(nd);
+        ++_size;
+    }
+
 public:
     constexpr iterator end(void) noexcept
     { return iterator(nullptr); }
@@ -583,39 +593,56 @@ public:
         return const_iterator { nextpos };
     }
 
+    template <typename U>
+    constexpr iterator lower_bound(U&& val) const
+    {
+        node* cur = root;
+        node* result = nullptr;
+
+        while (cur) {
+            if (!comp(cur->value, val)) {
+                result = cur;
+                cur = cur->left;
+            }
+            else {
+                cur = cur->right;
+            }
+        }
+
+        return iterator { result };
+    }
+
+    template <typename U>
+    constexpr iterator upper_bound(U&& val) const
+    {
+        iterator iter = lower_bound(std::forward<U>(val));
+        if (iter && !comp(*iter, val) && !comp(val, *iter))
+            return ++iter;
+        return iter;
+    }
+
     // value in nd MUST NOT exist in the rbtree,
     // that is, if a < b, then a > b
     constexpr void insert(node* nd)
     {
         node* cur = root;
 
-        while (cur) [[likely]] {
+        while (cur) {
             if (comp(nd->value, cur->value)) {
                 if (!cur->left) {
-                    nd->parent = cur;
-                    cur->left = nd;
-                    this->balance(nd);
-                    ++_size;
+                    do_insertion(cur, cur->left, nd);
                     return;
-                } else {
-                    cur = cur->left;
                 }
+                cur = cur->left;
             } else {
                 if (!cur->right) {
-                    nd->parent = cur;
-                    cur->right = nd;
-                    this->balance(nd);
-                    ++_size;
+                    do_insertion(cur, cur->right, nd);
                     return;
-                } else {
-                    cur = cur->right;
                 }
+                cur = cur->right;
             }
         }
-
-        root = nd;
-        root->toblack();
-        ++_size;
+        do_insertion(cur, root, nd);
     }
 
     template <typename U>

+ 14 - 0
gblibstdc++/include/map

@@ -245,6 +245,20 @@ public:
 
     __GBLIBCPP_CONSTEXPR
     bool contains(const Key& key) const { return count(key) != 0; }
+
+    __GBLIBCPP_CONSTEXPR
+    iterator upper_bound(const Key& key)
+    { return tree.upper_bound(key); }
+    __GBLIBCPP_CONSTEXPR
+    const_iterator upper_bound(const Key& key) const
+    { return tree.upper_bound(key); }
+
+    __GBLIBCPP_CONSTEXPR
+    iterator lower_bound(const Key& key)
+    { return tree.lower_bound(key); }
+    __GBLIBCPP_CONSTEXPR
+    const_iterator lower_bound(const Key& key) const
+    { return tree.lower_bound(key); }
 };
 
 template <typename Key, typename T, typename Compare, typename Allocator>

+ 14 - 0
gblibstdc++/include/set

@@ -166,6 +166,20 @@ public:
 
     __GBLIBCPP_CONSTEXPR
     bool contains(const Key& key) const { return count(key) != 0; }
+
+    __GBLIBCPP_CONSTEXPR
+    iterator upper_bound(const Key& key)
+    { return tree.upper_bound(key); }
+    __GBLIBCPP_CONSTEXPR
+    const_iterator upper_bound(const Key& key) const
+    { return tree.upper_bound(key); }
+
+    __GBLIBCPP_CONSTEXPR
+    iterator lower_bound(const Key& key)
+    { return tree.lower_bound(key); }
+    __GBLIBCPP_CONSTEXPR
+    const_iterator lower_bound(const Key& key) const
+    { return tree.lower_bound(key); }
 };
 
 template <typename Key, typename Compare, typename Allocator>