SHARE
TWEET

Untitled

a guest Apr 25th, 2019 65 Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
  1. #ifndef CPP_AVL_HPP
  2. #define CPP_AVL_HPP
  3.  
  4. #include <cstdint>
  5.  
  6. template<typename K>
  7. struct avl_node {
  8.   using node_t = avl_node<K>;
  9.  
  10.   node_t *cs_[2]{nullptr, nullptr};
  11.   node_t *p_{nullptr};
  12.   int8_t bf_{0};
  13.   uint8_t c_idx_{0};
  14.   K k_;
  15.  
  16.   explicit avl_node(K k) : k_(k) {}
  17.  
  18.   void set_child(int8_t c_idx, avl_node *child) {
  19.     cs_[c_idx] = child;
  20.     if (child != nullptr) {
  21.       child->p_ = this;
  22.       child->c_idx_ = c_idx;
  23.     }
  24.   }
  25.  
  26.   void set_child_nonnull(int8_t c_idx, avl_node *child) {
  27.     cs_[c_idx] = child;
  28.     child->p_ = this;
  29.     child->c_idx_ = c_idx;
  30.   }
  31. };
  32.  
  33. template<typename K, typename K_OP>
  34. struct avl_tree {
  35.   using node_t = avl_node<K>;
  36.  
  37.   node_t *root_;
  38.   size_t n_{0};
  39.  
  40.   avl_tree() : root_(nullptr) {}
  41.  
  42.   node_t *begin() const {
  43.     node_t *prev = nullptr;
  44.     for (node_t *node = root_; node != nullptr; node = node->cs_[0]) {
  45.       prev = node;
  46.     }
  47.     return prev;
  48.   }
  49.  
  50.   node_t *end() const { return nullptr; }
  51.  
  52.   node_t *rbegin() const {
  53.     node_t *prev = nullptr;
  54.     for (node_t *node = root_; node != nullptr; node = node->cs_[1]) {
  55.       prev = node;
  56.     }
  57.     return prev;
  58.   }
  59.  
  60.   node_t *rend() const { return nullptr; }
  61.  
  62.   node_t *next(node_t *node) const {
  63.     if (node->cs_[1] != nullptr) {
  64.       node = node->cs_[1];
  65.       while (node->cs_[0] != nullptr) node = node->cs_[0];
  66.     } else {
  67.       for (;;) {
  68.         uint8_t c_idx = node->c_idx_;
  69.         node = node->p_;
  70.         if (node == nullptr) return nullptr;
  71.         if (c_idx == 0) break;
  72.       }
  73.     }
  74.  
  75.     return node;
  76.   }
  77.  
  78.   node_t *prev(node_t *node) const {
  79.     if (node->cs_[0] != nullptr) {
  80.       node = node->cs_[0];
  81.       while (node->cs_[1] != nullptr) node = node->cs_[1];
  82.     } else {
  83.       for (;;) {
  84.         uint8_t c_idx = node->c_idx_;
  85.         node = node->p_;
  86.         if (node == nullptr) return nullptr;
  87.         if (c_idx == 1) break;
  88.       }
  89.     }
  90.  
  91.     return node;
  92.   }
  93.  
  94.   unsigned int rotate(node_t *node, node_t *parent, unsigned int reason) {
  95.     // return 0 if h(node) does not change
  96.     // return 1 if h(node) decrease by 1
  97.     int l = reason;
  98.     int r = 1 - l;
  99.     auto *P = node;
  100.     auto *C = P->cs_[l];
  101.     auto *N = C->cs_[r];
  102.     if (C->bf_ != 1 - 2 * l) {
  103.       /*           P
  104.        *     +-----+-----+
  105.        *     C
  106.        *  +--+--+
  107.        *        N
  108.        * becomes:
  109.        *           C
  110.        *     +-----+-----+
  111.        *                 P
  112.        *              +--+--+
  113.        *              N
  114.        *
  115.        * 1) l(0), P(-2), C( 0) -> P(-1), C(+1), return 0
  116.        * 2) l(0), P(-2), C(-1) -> P( 0), C( 0), return 1
  117.        * 1) l(1), P(+2), C( 0) -> P(+1), C(-1), return 0
  118.        * 2) l(1), P(+2), C(+1) -> P( 0), C( 0), return 1
  119.        */
  120.       if (parent == nullptr) root_ = C;
  121.       else parent->cs_[P->c_idx_] = C;
  122.       C->p_ = parent;
  123.       C->c_idx_ = P->c_idx_;
  124.  
  125.       C->set_child_nonnull(r, P);
  126.       P->set_child(l, N);
  127.  
  128.       C->bf_ += 1 - 2 * l;
  129.       P->bf_ = -C->bf_;
  130.  
  131.       return C->bf_ == 0;
  132.     } else {
  133.       /*          P
  134.        *    +-----+-----+
  135.        *    C
  136.        * +--+--+
  137.        *       N
  138.        *     +-+-+
  139.        *    Nl   Nr
  140.        * becomes:
  141.        *          N
  142.        *    +-----+-----+
  143.        *    C           P
  144.        * +--+--+     +--+--+
  145.        *      Nl     Nr
  146.        *
  147.        * 1) l(0), P(-2), C(+1), N(-1) -> P(+1), C( 0), N( 0), return 1
  148.        * 2) l(0), P(-2), C(+1), N( 0) -> P( 0), C( 0), N( 0), return 1
  149.        * 3) l(0), P(-2), C(+1), N(+1) -> P( 0), C(-1), N( 0), return 1
  150.        * 4) l(1), P(+2), C(-1), N(+1) -> P(-1), C( 0), N( 0), return 1
  151.        * 5) l(1), P(+2), C(-1), N( 0) -> P( 0), C( 0), N( 0), return 1
  152.        * 6) l(1), P(+2), C(-1), N(-1) -> P( 0), C(+1), N( 0), return 1
  153.        */
  154.       auto *Nl = N->cs_[l];
  155.       auto *Nr = N->cs_[r];
  156.  
  157.       if (parent == nullptr) root_ = N;
  158.       else parent->cs_[P->c_idx_] = N;
  159.       N->p_ = parent;
  160.       N->c_idx_ = P->c_idx_;
  161.  
  162.       N->set_child_nonnull(l, C);
  163.       N->set_child_nonnull(r, P);
  164.       C->set_child(r, Nl);
  165.       P->set_child(l, Nr);
  166.  
  167.       P->bf_ = (N->bf_ == 2 * l - 1) ? -N->bf_ : 0;
  168.       C->bf_ = (N->bf_ == 1 - 2 * l) ? -N->bf_ : 0;
  169.       N->bf_ = 0;
  170.  
  171.       return 1;
  172.     }
  173.   }
  174.  
  175.   avl_node<K> *find(const K k) {
  176.     if (root_ == nullptr) return nullptr;
  177.     auto *cursor = root_;
  178.     for (;;) {
  179.       if (K_OP::eq(k, cursor->k_)) return cursor;
  180.       int l = !K_OP::le(k, cursor->k_);
  181.       if (cursor->cs_[l] != nullptr) {
  182.         cursor = cursor->cs_[l];
  183.       } else {
  184.         return cursor;
  185.       }
  186.     }
  187.   }
  188.  
  189.   void insert(const K k) {
  190.     n_++;
  191.     auto *node = new node_t{k};
  192.     if (root_ == nullptr) {
  193.       root_ = node;
  194.     } else {
  195.       auto *cursor = root_;
  196.       for (;;) {
  197.         int l = !K_OP::le(k, cursor->k_);
  198.         if (cursor->cs_[l] != nullptr) {
  199.           cursor = cursor->cs_[l];
  200.         } else {
  201.           cursor->cs_[l] = node;
  202.           node->p_ = cursor;
  203.           node->c_idx_ = l;
  204.           cursor->bf_ += 2 * l - 1;
  205.           break;
  206.         }
  207.       }
  208.  
  209.       while (!(cursor->bf_ == 0 || cursor->p_ == nullptr)) {
  210.         if (cursor->p_->bf_ == 2 * cursor->c_idx_ - 1) {
  211.           rotate(cursor->p_, cursor->p_->p_, cursor->c_idx_);
  212.           break;
  213.         }
  214.         cursor->p_->bf_ += 2 * cursor->c_idx_ - 1;
  215.         cursor = cursor->p_;
  216.       }
  217.     }
  218.   }
  219.  
  220.   void erase(K k) {
  221.     n_--;
  222.     auto *target = root_;
  223.     while (!K_OP::eq(k, target->k_)) target = target->cs_[!K_OP::le(k, target->k_)];
  224.  
  225.     node_t *to_delete = target;
  226.     unsigned int l = (target->bf_ + 1) / 2;
  227.     if (target->cs_[0] != nullptr && target->cs_[1] != nullptr) {
  228.       int r = 1 - l;
  229.       to_delete = target->cs_[l];
  230.       while (to_delete->cs_[r] != nullptr) to_delete = to_delete->cs_[r];
  231.     }
  232.  
  233.     target->k_ = std::move(to_delete->k_);
  234.  
  235.     if (to_delete->p_ == nullptr) {
  236.       root_ = to_delete->cs_[l];
  237.       if (root_ != nullptr) root_->p_ = nullptr;
  238.       delete to_delete;
  239.       return;
  240.     }
  241.     auto *cursor = to_delete->p_;
  242.     auto c_idx = to_delete->c_idx_;
  243.     cursor->set_child(c_idx, to_delete->cs_[l]);
  244.     delete to_delete;
  245.  
  246.     l = !c_idx;
  247.     while (cursor) {
  248.       auto *p = cursor->p_;
  249.       auto idx = cursor->c_idx_;
  250.       if (cursor->bf_ == 2 * l - 1) {
  251.         if (!rotate(cursor, p, l)) return;
  252.       } else {
  253.         cursor->bf_ += 2 * l - 1;
  254.         if (cursor->bf_ != 0) return;
  255.       }
  256.       cursor = p;
  257.       l = !idx;
  258.     }
  259.   }
  260. };
  261.  
  262. #endif //CPP_AVL_HPP
RAW Paste Data
We use cookies for various purposes including analytics. By continuing to use Pastebin, you agree to our use of cookies as described in the Cookies Policy. OK, I Understand
 
Top