tpl_rb_tree.H

Go to the documentation of this file.

 
/*
                          Aleph_w
 
  Data structures & Algorithms
  version 2.0.0b
  https://github.com/lrleon/Aleph-w
 
  This file is part of Aleph-w library
 
  Copyright (c) 2002-2026 Leandro Rabindranath Leon
 
  Permission is hereby granted, free of charge, to any person obtaining a copy
  of this software and associated documentation files (the "Software"), to deal
  in the Software without restriction, including without limitation the rights
  to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  copies of the Software, and to permit persons to whom the Software is
  furnished to do so, subject to the following conditions:
 
  The above copyright notice and this permission notice shall be included in all
  copies or substantial portions of the Software.
 
  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  SOFTWARE.
*/
 
 
# ifndef TPL_RB_TREE_H
# define TPL_RB_TREE_H
 
# include <ahFunction.H>
# include <tpl_arrayStack.H>
# include <tpl_binNodeUtils.H>
# include <rbNode.H>
# include <ah-concepts.H>
 
namespace Aleph
{
 
template <template <typename> class NodeType, typename Key, class Compare>
    requires StrictWeakOrder<Compare, Key>
class Gen_Rb_Tree
{
public:
  using Node = NodeType<Key>;
  using key_type = Key;
  using compare_type = Compare;
 
private:
  Node head_node;                    
  Node *head;                        
  Node *& root;                      
  FixedStack<Node *> rb_stack;       
  Compare cmp;                       
  size_t num_nodes;                  
  static constexpr signed char CmpLess = -1;
  static constexpr signed char CmpEqual = 0;
  static constexpr signed char CmpGreater = 1;
 
  Node * search_and_stack_rb(const Key & key,
                             signed char & cmp_result) noexcept
  {
    Node *p = root;
    Node *candidate = nullptr;  // Tracks potential duplicate when going right
    size_t candidate_pos = 0;   // Stack size when candidate was set
 
    rb_stack.push(head);
    do
      {
        rb_stack.push(p);
        if (cmp(key, KEY(p)))  // key < KEY(p)
          {
            cmp_result = CmpLess;
            p = LLINK(p);
          }
        else  // key >= KEY(p), could be equal
          {
            candidate = p;
            candidate_pos = rb_stack.size();
            cmp_result = CmpGreater;
            p = RLINK(p);
          }
      }
    while (p != Node::NullPtr);
 
    // Optimistic duplicate check: when going right, key >= KEY(candidate)
    // If also KEY(candidate) >= key (i.e., not less), then key == KEY(candidate)
    if (candidate != nullptr and not cmp(KEY(candidate), key)) [[unlikely]]
      {
        cmp_result = CmpEqual;
        // Truncate stack: remove nodes pushed after candidate
        const size_t to_pop =
            rb_stack.size() > candidate_pos ?
            rb_stack.size() - candidate_pos : 0;
        if (to_pop > 0)
          rb_stack.popn(static_cast<int>(to_pop));
        return candidate;
      }
 
    return rb_stack.top();
  }
 
  Node * search_dup_and_stack_rb(const Key & key,
                                 signed char & cmp_result) noexcept
  {
    Node *p = root;
    rb_stack.push(head);
    cmp_result = CmpEqual;
    do
      {
        rb_stack.push(p);
        const Key & pk = KEY(p);
        if (cmp(key, pk))
          {
            cmp_result = CmpLess;
            p = LLINK(p);
          }
        else
          {
            cmp_result = CmpGreater;
            p = RLINK(p);
          }
      }
    while (p != Node::NullPtr);
 
    return rb_stack.top();
  }
 
  void fix_red_condition(Node *p) noexcept
  {
    assert(COLOR(p) == RED);
 
    while (p != root)
      {
        Node *pp = rb_stack.pop();  // parent of p
        if (COLOR(pp) == BLACK)     // is p's parent black?
          break;                    // yes ==> no consecutive reds ==> done
 
        if (root == pp)             // is p a direct child of the root?
          {
            COLOR(root) = BLACK;
            break;
          }
 
        Node *ppp = rb_stack.pop();  // grandparent of p
        Node *spp = LLINK(ppp) == pp ? RLINK(ppp) : LLINK(ppp);  // uncle
        
        if (COLOR(spp) == RED)       // is p's uncle red?
          {
            COLOR(ppp) = RED;
            COLOR(pp) = BLACK;
            COLOR(spp) = BLACK;
            p = ppp;
            continue;
          }
 
        Node *pppp = rb_stack.pop();  // great-grandparent of p
        if (LLINK(pp) == p and LLINK(ppp) == pp)
          {
            rotate_to_right(ppp, pppp);
            COLOR(pp) = BLACK;
          }
        else if (RLINK(pp) == p and RLINK(ppp) == pp)
          {
            rotate_to_left(ppp, pppp);
            COLOR(pp) = BLACK;
          }
        else
          {
            if (RLINK(pp) == p)
              {
                rotate_to_left(pp, ppp);
                rotate_to_right(ppp, pppp);
              }
            else
              {
                rotate_to_right(pp, ppp);
                rotate_to_left(ppp, pppp);
              }
            COLOR(p) = BLACK;
          }
 
        COLOR(ppp) = RED;
        break;
      }
 
    rb_stack.empty();
  }
 
  void find_succ_and_swap(Node *p, Node *& pp) noexcept
  {
    Node *& ref_rb_stack = rb_stack.top();
 
    Node *fSucc = p;
    Node *succ = RLINK(p);
    rb_stack.push(succ);
 
    while (LLINK(succ) != Node::NullPtr)
      {
        fSucc = succ;
        succ = LLINK(succ);
        rb_stack.push(succ);
      }
 
    ref_rb_stack = succ;
    rb_stack.top() = p;
 
    if (LLINK(pp) == p)
      LLINK(pp) = succ;
    else
      RLINK(pp) = succ;
 
    LLINK(succ) = LLINK(p);
    LLINK(p) = Node::NullPtr;
 
    if (RLINK(p) == succ)
      {
        RLINK(p) = RLINK(succ);
        RLINK(succ) = p;
        pp = succ;
      }
    else
      {
        Node *succr = RLINK(succ);
        RLINK(succ) = RLINK(p);
        LLINK(fSucc) = p;
        RLINK(p) = succr;
        pp = fSucc;
      }
 
    std::swap(COLOR(succ), COLOR(p));
  }
 
  void fix_black_condition(Node *p) noexcept
  {
    if (COLOR(p) == RED)
      {
        COLOR(p) = BLACK;
        return;
      }
 
    Node *pp = rb_stack.popn(2);
    while (p != root)
      {
        assert(LLINK(pp) == p or RLINK(pp) == p);
        assert(LLINK(rb_stack.top()) == pp or RLINK(rb_stack.top()) == pp);
 
        Node *sp = LLINK(pp) == p ? RLINK(pp) : LLINK(pp);
        if (COLOR(sp) == RED)
          {
            Node *& ppp = rb_stack.top();
 
            if (LLINK(pp) == p)
              {
                sp = LLINK(sp);
                ppp = rotate_to_left(pp, ppp);
              }
            else
              {
                sp = RLINK(sp);
                ppp = rotate_to_right(pp, ppp);
              }
 
            COLOR(ppp) = BLACK;
            COLOR(pp) = RED;
          }
 
        Node *np, *snp;
        if (LLINK(pp) == p)
          {
            np = RLINK(sp);
            snp = LLINK(sp);
          }
        else
          {
            np = LLINK(sp);
            snp = RLINK(sp);
          }
 
        if (COLOR(np) == RED)
          {
            Node *ppp = rb_stack.top();
 
            if (RLINK(sp) == np)
              rotate_to_left(pp, ppp);
            else
              rotate_to_right(pp, ppp);
 
            COLOR(sp) = COLOR(pp);
            COLOR(pp) = BLACK;
            COLOR(np) = BLACK;
 
            return;
          }
 
        if (COLOR(snp) == RED)
          {
            Node *ppp = rb_stack.top();
 
            if (LLINK(sp) == snp)
              {
                rotate_to_right(sp, pp);
                rotate_to_left(pp, ppp);
              }
            else
              {
                rotate_to_left(sp, pp);
                rotate_to_right(pp, ppp);
              }
 
            COLOR(snp) = COLOR(pp);
            COLOR(pp) = BLACK;
 
            return;
          }
 
        if (COLOR(pp) == RED)
          {
            COLOR(pp) = BLACK;
            COLOR(sp) = RED;
            return;
          }
 
        COLOR(sp) = RED;
        p = pp;
        pp = rb_stack.pop();
      }
  }
 
  void init() noexcept
  {
    num_nodes = 0;
  }
 
public:
  Compare & key_comp() noexcept { return cmp; }
  const Compare & key_comp() const noexcept { return cmp; }
 
  Compare & get_compare() noexcept { return cmp; }
  const Compare & get_compare() const noexcept { return cmp; }
 
  Gen_Rb_Tree(Compare __cmp = Compare()) noexcept
    : head(&head_node), root(head_node.getR()),
      rb_stack(Node::MaxHeight), cmp(__cmp), num_nodes(0)
  {
    /* empty */
  }
 
  void swap(Gen_Rb_Tree & tree) noexcept
  {
    std::swap(root, tree.root);
    std::swap(num_nodes, tree.num_nodes);
    std::swap(cmp, tree.cmp);
  }
 
  Gen_Rb_Tree(Gen_Rb_Tree && tree) noexcept
    : head(&head_node), root(head_node.getR()),
      rb_stack(Node::MaxHeight), cmp(std::move(tree.cmp)), 
      num_nodes(tree.num_nodes)
  {
    root = tree.root;
    tree.root = Node::NullPtr;
    tree.num_nodes = 0;
  }
 
  Gen_Rb_Tree & operator=(Gen_Rb_Tree && tree) noexcept
  {
    if (this != &tree)
      {
        root = tree.root;
        num_nodes = tree.num_nodes;
        cmp = std::move(tree.cmp);
        tree.root = Node::NullPtr;
        tree.num_nodes = 0;
      }
    return *this;
  }
 
  virtual ~Gen_Rb_Tree() = default;
 
  Node * search(const Key & key) const noexcept
  {
    Node *p = root;
    while (p != Node::NullPtr)
      {
        if (cmp(key, KEY(p)))
          p = LLINK(p);
        else if (cmp(KEY(p), key))
          p = RLINK(p);
        else
          return p;
      }
    return nullptr;
  }
 
  Node *& getRoot() noexcept { return root; }
  
  Node * getRoot() const noexcept { return root; }
 
  bool is_empty() const noexcept { return root == Node::NullPtr; }
 
  size_t size() const noexcept { return num_nodes; }
 
  void reset() noexcept
  {
    root = Node::NullPtr;
    num_nodes = 0;
  }
 
  Node * insert(Node *p) noexcept
  {
    assert(p != nullptr and p != Node::NullPtr);
    assert(COLOR(p) == RED);
 
    if (root == Node::NullPtr)
      {
        ++num_nodes;
        return root = p;
      }
 
    signed char cmp_result = CmpEqual;
    Node *q = search_and_stack_rb(KEY(p), cmp_result);
    if (cmp_result == CmpLess)
      LLINK(q) = p;
    else if (cmp_result == CmpGreater)
      RLINK(q) = p;
    else
      {
        rb_stack.empty();
        return nullptr;  // duplicate key
      }
    
    ++num_nodes;
    fix_red_condition(p);
 
    return p;
  }
 
  Node * search_or_insert(Node *p) noexcept
  {
    assert(p != nullptr and p != Node::NullPtr);
    assert(COLOR(p) == RED);
 
    if (root == Node::NullPtr)
      {
        ++num_nodes;
        return root = p;
      }
 
    signed char cmp_result = CmpEqual;
    Node *q = search_and_stack_rb(KEY(p), cmp_result);
    if (cmp_result == CmpLess)
      LLINK(q) = p;
    else if (cmp_result == CmpGreater)
      RLINK(q) = p;
    else
      {
        rb_stack.empty();
        return q;  // return existing node
      }
    
    ++num_nodes;
    fix_red_condition(p);
 
    return p;
  }
 
  Node * insert_dup(Node *p) noexcept
  {
    assert(p != nullptr and p != Node::NullPtr);
    assert(COLOR(p) == RED);
 
    if (root == Node::NullPtr)
      {
        ++num_nodes;
        return root = p;
      }
 
    signed char cmp_result = CmpEqual;
    Node *q = search_dup_and_stack_rb(KEY(p), cmp_result);
    if (cmp_result == CmpLess)
      LLINK(q) = p;
    else
      RLINK(q) = p;
 
    ++num_nodes;
    fix_red_condition(p);
 
    return p;
  }
 
  bool verify() const noexcept { return is_red_black_bst(root, cmp); }
 
  Node * remove(const Key & key) noexcept
  {
    if (root == Node::NullPtr)
      return nullptr;
 
    signed char cmp_result = CmpEqual;
    Node *q = search_and_stack_rb(key, cmp_result);
    if (cmp_result != CmpEqual)
      {
        rb_stack.empty();
        return nullptr;
      }
 
    Node *pq = rb_stack.top(1);
    Node *p;
    while (true)
      {
        if (LLINK(q) == Node::NullPtr)
          {
            if (LLINK(pq) == q)
              p = LLINK(pq) = RLINK(q);
            else
              p = RLINK(pq) = RLINK(q);
            break;
          }
 
        if (RLINK(q) == Node::NullPtr)
          {
            if (LLINK(pq) == q)
              p = LLINK(pq) = LLINK(q);
            else
              p = RLINK(pq) = LLINK(q);
            break;
          }
 
        find_succ_and_swap(q, pq);
      }
 
    if (COLOR(q) == BLACK)
      fix_black_condition(p);
 
    q->reset();
    rb_stack.empty();
    --num_nodes;
 
    return q;
  }
 
  struct Iterator : public BinNodeInfixIterator<Node>
  {
    Iterator() noexcept = default;
    Iterator(Gen_Rb_Tree & t) : BinNodeInfixIterator<Node>(t.getRoot()) {}
    Iterator(const Gen_Rb_Tree & t) : BinNodeInfixIterator<Node>(t.getRoot()) {}
  };
};
 
template <typename Key, class Compare = Aleph::less<Key>>
    requires StrictWeakOrder<Compare, Key>
struct Rb_Tree : public Gen_Rb_Tree<RbNode, Key, Compare>
{
  using Base = Gen_Rb_Tree<RbNode, Key, Compare>;
  using Base::Base;
};
 
template <typename Key, class Compare = Aleph::less<Key>>
    requires StrictWeakOrder<Compare, Key>
struct Rb_Tree_Vtl : public Gen_Rb_Tree<RbNodeVtl, Key, Compare>
{
  using Base = Gen_Rb_Tree<RbNodeVtl, Key, Compare>;
  using Base::Base;
};
 
} // end namespace Aleph
 
# endif // TPL_RB_TREE_H