Advertisement
Not a member of Pastebin yet?
Sign Up,
it unlocks many cool features!
- #ifndef CPP_AVL_HPP
- #define CPP_AVL_HPP
- #include <cstdint>
- template<typename K>
- struct avl_node {
- using node_t = avl_node<K>;
- node_t *cs_[2]{nullptr, nullptr};
- node_t *p_{nullptr};
- int8_t bf_{0};
- uint8_t c_idx_{0};
- K k_;
- explicit avl_node(K k) : k_(k) {}
- void set_child(int8_t c_idx, avl_node *child) {
- cs_[c_idx] = child;
- if (child != nullptr) {
- child->p_ = this;
- child->c_idx_ = c_idx;
- }
- }
- void set_child_nonnull(int8_t c_idx, avl_node *child) {
- cs_[c_idx] = child;
- child->p_ = this;
- child->c_idx_ = c_idx;
- }
- };
- template<typename K, typename K_OP>
- struct avl_tree {
- using node_t = avl_node<K>;
- node_t *root_;
- size_t n_{0};
- avl_tree() : root_(nullptr) {}
- node_t *begin() const {
- node_t *prev = nullptr;
- for (node_t *node = root_; node != nullptr; node = node->cs_[0]) {
- prev = node;
- }
- return prev;
- }
- node_t *end() const { return nullptr; }
- node_t *rbegin() const {
- node_t *prev = nullptr;
- for (node_t *node = root_; node != nullptr; node = node->cs_[1]) {
- prev = node;
- }
- return prev;
- }
- node_t *rend() const { return nullptr; }
- node_t *next(node_t *node) const {
- if (node->cs_[1] != nullptr) {
- node = node->cs_[1];
- while (node->cs_[0] != nullptr) node = node->cs_[0];
- } else {
- for (;;) {
- uint8_t c_idx = node->c_idx_;
- node = node->p_;
- if (node == nullptr) return nullptr;
- if (c_idx == 0) break;
- }
- }
- return node;
- }
- node_t *prev(node_t *node) const {
- if (node->cs_[0] != nullptr) {
- node = node->cs_[0];
- while (node->cs_[1] != nullptr) node = node->cs_[1];
- } else {
- for (;;) {
- uint8_t c_idx = node->c_idx_;
- node = node->p_;
- if (node == nullptr) return nullptr;
- if (c_idx == 1) break;
- }
- }
- return node;
- }
- unsigned int rotate(node_t *node, node_t *parent, unsigned int reason) {
- // return 0 if h(node) does not change
- // return 1 if h(node) decrease by 1
- int l = reason;
- int r = 1 - l;
- auto *P = node;
- auto *C = P->cs_[l];
- auto *N = C->cs_[r];
- if (C->bf_ != 1 - 2 * l) {
- /* P
- * +-----+-----+
- * C
- * +--+--+
- * N
- * becomes:
- * C
- * +-----+-----+
- * P
- * +--+--+
- * N
- *
- * 1) l(0), P(-2), C( 0) -> P(-1), C(+1), return 0
- * 2) l(0), P(-2), C(-1) -> P( 0), C( 0), return 1
- * 1) l(1), P(+2), C( 0) -> P(+1), C(-1), return 0
- * 2) l(1), P(+2), C(+1) -> P( 0), C( 0), return 1
- */
- if (parent == nullptr) root_ = C;
- else parent->cs_[P->c_idx_] = C;
- C->p_ = parent;
- C->c_idx_ = P->c_idx_;
- C->set_child_nonnull(r, P);
- P->set_child(l, N);
- C->bf_ += 1 - 2 * l;
- P->bf_ = -C->bf_;
- return C->bf_ == 0;
- } else {
- /* P
- * +-----+-----+
- * C
- * +--+--+
- * N
- * +-+-+
- * Nl Nr
- * becomes:
- * N
- * +-----+-----+
- * C P
- * +--+--+ +--+--+
- * Nl Nr
- *
- * 1) l(0), P(-2), C(+1), N(-1) -> P(+1), C( 0), N( 0), return 1
- * 2) l(0), P(-2), C(+1), N( 0) -> P( 0), C( 0), N( 0), return 1
- * 3) l(0), P(-2), C(+1), N(+1) -> P( 0), C(-1), N( 0), return 1
- * 4) l(1), P(+2), C(-1), N(+1) -> P(-1), C( 0), N( 0), return 1
- * 5) l(1), P(+2), C(-1), N( 0) -> P( 0), C( 0), N( 0), return 1
- * 6) l(1), P(+2), C(-1), N(-1) -> P( 0), C(+1), N( 0), return 1
- */
- auto *Nl = N->cs_[l];
- auto *Nr = N->cs_[r];
- if (parent == nullptr) root_ = N;
- else parent->cs_[P->c_idx_] = N;
- N->p_ = parent;
- N->c_idx_ = P->c_idx_;
- N->set_child_nonnull(l, C);
- N->set_child_nonnull(r, P);
- C->set_child(r, Nl);
- P->set_child(l, Nr);
- P->bf_ = (N->bf_ == 2 * l - 1) ? -N->bf_ : 0;
- C->bf_ = (N->bf_ == 1 - 2 * l) ? -N->bf_ : 0;
- N->bf_ = 0;
- return 1;
- }
- }
- avl_node<K> *find(const K k) {
- if (root_ == nullptr) return nullptr;
- auto *cursor = root_;
- for (;;) {
- if (K_OP::eq(k, cursor->k_)) return cursor;
- int l = !K_OP::le(k, cursor->k_);
- if (cursor->cs_[l] != nullptr) {
- cursor = cursor->cs_[l];
- } else {
- return cursor;
- }
- }
- }
- void insert(const K k) {
- n_++;
- auto *node = new node_t{k};
- if (root_ == nullptr) {
- root_ = node;
- } else {
- auto *cursor = root_;
- for (;;) {
- int l = !K_OP::le(k, cursor->k_);
- if (cursor->cs_[l] != nullptr) {
- cursor = cursor->cs_[l];
- } else {
- cursor->cs_[l] = node;
- node->p_ = cursor;
- node->c_idx_ = l;
- cursor->bf_ += 2 * l - 1;
- break;
- }
- }
- while (!(cursor->bf_ == 0 || cursor->p_ == nullptr)) {
- if (cursor->p_->bf_ == 2 * cursor->c_idx_ - 1) {
- rotate(cursor->p_, cursor->p_->p_, cursor->c_idx_);
- break;
- }
- cursor->p_->bf_ += 2 * cursor->c_idx_ - 1;
- cursor = cursor->p_;
- }
- }
- }
- void erase(K k) {
- n_--;
- auto *target = root_;
- while (!K_OP::eq(k, target->k_)) target = target->cs_[!K_OP::le(k, target->k_)];
- node_t *to_delete = target;
- unsigned int l = (target->bf_ + 1) / 2;
- if (target->cs_[0] != nullptr && target->cs_[1] != nullptr) {
- int r = 1 - l;
- to_delete = target->cs_[l];
- while (to_delete->cs_[r] != nullptr) to_delete = to_delete->cs_[r];
- }
- target->k_ = std::move(to_delete->k_);
- if (to_delete->p_ == nullptr) {
- root_ = to_delete->cs_[l];
- if (root_ != nullptr) root_->p_ = nullptr;
- delete to_delete;
- return;
- }
- auto *cursor = to_delete->p_;
- auto c_idx = to_delete->c_idx_;
- cursor->set_child(c_idx, to_delete->cs_[l]);
- delete to_delete;
- l = !c_idx;
- while (cursor) {
- auto *p = cursor->p_;
- auto idx = cursor->c_idx_;
- if (cursor->bf_ == 2 * l - 1) {
- if (!rotate(cursor, p, l)) return;
- } else {
- cursor->bf_ += 2 * l - 1;
- if (cursor->bf_ != 0) return;
- }
- cursor = p;
- l = !idx;
- }
- }
- };
- #endif //CPP_AVL_HPP
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement