AVL Tree

Cây AVL là một loại cây tìm kiếm nhị phân cân bằng. Do cách giới thiệu về AVL trong nhiều giáo trình thuật toán thường rất dài dòng, khiến nhiều người có ấn tượng rằng cây AVL phức tạp và không thực dụng. Tuy nhiên, trên thực tế, nguyên lý của cây AVL rất đơn giản và việc cài đặt cũng không hề phức tạp.

Tính chất

1. Cây nhị phân rỗng là một cây AVL.
2. Nếu T là một cây AVL, thì các cây con bên trái và bên phải của nó cũng là cây AVL, và \(|h(ls) - h(rs)| \leq 1\), trong đó \(h\) là chiều cao của các cây con.
3. Chiều cao của cây là \(O(\log n)\).

Hệ số cân bằng: Chiều cao cây con bên phải - Chiều cao cây con bên trái.

Chứng minh chiều cao cây

Gọi \(f_n\) là số nút tối thiểu mà một cây AVL có chiều cao \(n\) chứa, ta có:

\[ f_n= \begin{cases} 1&(n=1)\\ 2&(n=2)\\ f_{n-1}+f_{n-2}+1& (n>2) \end{cases} \]

Theo cách giải phương trình sai phân tuyến tính không thuần nhất với hệ số hằng, \(\{f_n+1\}\) là một dãy Fibonacci. Công thức tổng quát của \(f_n\) là:

\[ f_n=\frac{5+2\sqrt{5}}{5}\left(\frac{1+\sqrt{5}}{2}\right)^n+\frac{5-2\sqrt{5}}{5}\left(\frac{1-\sqrt{5}}{2}\right)^n-1 \]

Dãy Fibonacci tăng trưởng theo tốc độ hàm mũ, đối với chiều cao cây \(n\) ta có:

\[ n<\log_{\frac{1+\sqrt{5}}{2}} (f_n+1)<\frac{3}{2}\log_2 (f_n+1) \]

Do đó chiều cao của cây AVL là \(O(\log f_n)\), với \(f_n\) là số lượng nút.

Quá trình

Chèn nút

Tương tự như trong BST (Cây tìm kiếm nhị phân), trước tiên thực hiện một lần tìm kiếm thất bại để xác định vị trí chèn. Sau khi chèn nút, dựa vào hệ số cân bằng để quyết định có cần điều chỉnh hay không.

Xóa nút

Xóa nút tương tự như BST, đổi chỗ nút cần xóa với nút kế nhiệm (successor) rồi tiến hành xóa.

Việc xóa có thể dẫn đến thay đổi chiều cao cây và hệ số cân bằng, lúc này cần điều chỉnh dọc theo con đường từ nút bị xóa lên đến gốc.

Duy trì sự cân bằng

Sau khi chèn hoặc xóa nút, tính chất 2 của cây AVL có thể bị phá vỡ. Do đó, cần duy trì cây dọc theo con đường từ nút vừa chèn/xóa đến gốc. Nếu tại một nút nào đó tính chất 2 không còn được thỏa mãn, vì chúng ta chỉ chèn/xóa một nút nên ảnh hưởng đến chiều cao cây không quá 1, do đó giá trị tuyệt đối của hệ số cân bằng tại nút đó tối đa là 2. Do tính đối xứng, ở đây chúng ta chỉ thảo luận trường hợp chiều cao cây con bên trái lớn hơn cây con bên phải là 2, tức là trong hình dưới đây \(h(B)-h(E)=2\). Lúc này, cần thảo luận hai trường hợp dựa trên mối quan hệ giữa \(h(A)\) và \(h(C)\). Cần lưu ý rằng, vì chúng ta duy trì sự cân bằng từ dưới lên trên, nên đối với tất cả các hậu duệ của nút D, tính chất 2 vẫn được thỏa mãn.

Trường hợp 1: Chiều cao nút A không nhỏ hơn chiều cao nút C

Gọi \(h(E)=x\), ta có:

\[ \begin{cases} h(B)=x+2\\ h(A)=x+1\\ x\leq h(C)\leq x+1 \end{cases} \]

Trong đó \(h(C)\geq x\) là vì nút B thỏa mãn tính chất 2, nên chênh lệch giữa \(h(C)\) và \(h(A)\) không quá 1. Lúc này chúng ta thực hiện một thao tác xoay phải tại nút D (thao tác xoay giống như các loại cây tìm kiếm nhị phân cân bằng khác), như hình dưới đây.

Rõ ràng chiều cao các nút A, C, E không thay đổi, và ta có:

\[ \begin{cases} 0\leq h(C)-h(E)\leq 1\\ x+1\leq h'(D)=\max(h(C),h(E))+1=h(C)+1\leq x+2\\ 0\leq h'(D)-h(A)\leq 1 end{cases} \]

Vì vậy sau khi xoay, các nút B và D cũng thỏa mãn tính chất 2.

Trường hợp 2: Chiều cao nút A nhỏ hơn chiều cao nút C

Gọi \(h(E)=x\), tương tự như trên, ta có:

\[ \begin{cases} h(B)=x+2\\ h(C)=x+1\\ h(A)=x \end{cases} \]

Lúc này trước tiên thực hiện một thao tác xoay trái tại nút B, sau đó thực hiện một thao tác xoay phải tại nút D, như hình dưới đây.

Rõ ràng chiều cao các nút A, E không thay đổi, và con bên phải mới của B cùng con bên trái mới của D lần lượt là con trái và con phải ban đầu của C, ta có:

\[ \begin{cases} x-1\leq h'(rs_B),h'(ls_D)\leq x\\ 0\leq h(A)-h'(rs_B)\leq 1\\ 0\leq h(E)-h'(ls_D)\leq 1\\ h'(B)=\max(h(A),h'(rs_B))+1=x+1\\ h'(D)=\max(h(E),h'(ls_D))+1=x+1\\ h'(B)-h'(D)=0 \end{cases} \]

Vì vậy sau khi xoay, các nút B, C, D cũng thỏa mãn tính chất 2.

Thao tác duy trì cân bằng: Mã giả

\[ \begin{array}{ll} 1 & \textbf{function } \mathrm{MaintainBalance}(p) \\ 2 & \qquad l \gets ls_p, r \gets rs_p \\ 3 & \qquad \textbf{if } h(l)-h(r)=2 \\ 4 & \qquad\qquad \textbf{if } h(ls_l) \ge h(rs_l) \\ 5 & \qquad\qquad\qquad \mathrm{RightRotate}(p) \\ 6 & \qquad\qquad \textbf{else} \\ 7 & \qquad\qquad\qquad \mathrm{LeftRotate}(l) \\ 8 & \qquad\qquad\qquad \mathrm{RightRotate}(p) \\ 9 & \qquad \textbf{else if } h(l)-h(r)=-2 \\ 10 & \qquad\qquad \textbf{if } h(ls_r) \le h(rs_r) \\ 11 & \qquad\qquad\qquad \mathrm{LeftRotate}(p) \\ 12 & \qquad\qquad \textbf{else} \\ 13 & \qquad\qquad\qquad \mathrm{RightRotate}(r) \\ 14 & \qquad\qquad\qquad \mathrm{LeftRotate}(p) \\ \end{array} \]

Giống như các cây tìm kiếm nhị phân cân bằng khác, các thông tin như chiều cao nút, kích thước cây con của cây AVL cần được duy trì khi xoay.

Các thao tác khác

Các thao tác khác của cây AVL (Predecessor, Successor, Select, Rank, v.v.) giống như cây tìm kiếm nhị phân thông thường.

Mã tham khảo

Mã dưới đây là một Map được cài đặt bằng cây AVL (ánh xạ có thứ tự không lặp lại):

Mã tham khảo

/**
 * @brief An AVLTree-based map implementation
 * @details The map is sorted according to the natural ordering of its
 *  keys or by a {@code Compare} function provided; This implementation
 *  provides guaranteed log(n) time cost for the contains, get, insert
 *  and remove operations.
 */

#ifndef AVLTREE_MAP_HPP
#define AVLTREE_MAP_HPP

#include <cassert>
#include <cstddef>
#include <cstdint>
#include <functional>
#include <memory>
#include <stack>
#include <utility>
#include <vector>

/**
 * An AVLTree-based map implementation
 * https://en.wikipedia.org/wiki/AVL_tree
 * @tparam Key the type of keys maintained by this map
 * @tparam Value the type of mapped values
 * @tparam Compare
 */
template <typename Key, typename Value, typename Compare = std::less<Key> >
class AvlTreeMap {
 private:
  using USize = size_t;
  using Factor = int64_t;

  Compare compare = Compare();

 public:
  struct Entry {
    Key key;
    Value value;

    bool operator==(const Entry &rhs) const noexcept {
      return this->key == rhs.key && this->value == rhs.value;
    }

    bool operator!=(const Entry &rhs) const noexcept {
      return this->key != rhs.key || this->value != rhs.value;
    }
  };

 private:
  struct Node {
    using Ptr = std::shared_ptr<Node>;
    using Provider = const std::function<Ptr(void)> &;
    using Consumer = const std::function<void(const Ptr &)> &;

    Key key;
    Value value{};

    Ptr left = nullptr;
    Ptr right = nullptr;

    USize height = 1;

    explicit Node(Key k) : key(std::move(k)) {}

    explicit Node(Key k, Value v) : key(std::move(k)), value(std::move(v)) {}

    ~Node() = default;

    inline bool isLeaf() const noexcept {
      return this->left == nullptr && this->right == nullptr;
    }

    inline void updateHeight() noexcept {
      if (this->isLeaf()) {
        this->height = 1;
      } else if (this->left == nullptr) {
        this->height = this->right->height + 1;
      } else if (this->right == nullptr) {
        this->height = this->left->height + 1;
      } else {
        this->height = std::max(left->height, right->height) + 1;
      }
    }

    inline Factor factor() const noexcept {
      if (this->isLeaf()) {
        return 0;
      } else if (this->left == nullptr) {
        return (Factor)this->right->height;
      } else if (this->right == nullptr) {
        return (Factor) - this->left->height;
      } else {
        return (Factor)(this->right->height - this->left->height);
      }
    }

    inline Entry entry() const { return Entry{key, value}; }

    static Ptr from(const Key &k) { return std::make_shared<Node>(Node(k)); }

    static Ptr from(const Key &k, const Value &v) {
      return std::make_shared<Node>(Node(k, v));
    }
  };

  using NodePtr = typename Node::Ptr;
  using ConstNodePtr = const NodePtr &;
  using NodeProvider = typename Node::Provider;
  using NodeConsumer = typename Node::Consumer;

  NodePtr root = nullptr;
  USize count = 0;

  using K = const Key &;
  using V = const Value &;

 public:
  using EntryList = std::vector<Entry>;
  using KeyValueConsumer = const std::function<void(K, V)> &;
  using MutKeyValueConsumer = const std::function<void(K, Value &)> &;
  using KeyValueFilter = const std::function<bool(K, V)> &;

  class NoSuchMappingException : protected std::exception {
   private:
    const char *message;

   public:
    explicit NoSuchMappingException(const char *msg) : message(msg) {}

    const char *what() const noexcept override { return message; }
  };

  AvlTreeMap() noexcept = default;

  /**
   * Returns the number of entries in this map.
   * @return size_t
   */
  inline USize size() const noexcept { return this->count; }

  /**
   * Returns true if this collection contains no elements.
   * @return bool
   */
  inline bool empty() const noexcept { return this->count == 0; }

  /**
   * Removes all of the elements from this map.
   */
  void clear() noexcept {
    this->root = nullptr;
    this->count = 0;
  }

  /**
   * Returns the value to which the specified key is mapped; If this map
   * contains no mapping for the key, a {@code NoSuchMappingException} will
   * be thrown.
   * @param key
   * @return AvlTreeMap<Key, Value>::Value
   * @throws NoSuchMappingException
   */
  Value get(K key) const {
    if (this->root == nullptr) {
      throw NoSuchMappingException("Invalid key");
    } else {
      NodePtr node = this->getNode(this->root, key);
      if (node != nullptr) {
        return node->value;
      } else {
        throw NoSuchMappingException("Invalid key");
      }
    }
  }

  /**
   * Returns the value to which the specified key is mapped; If this map
   * contains no mapping for the key, a new mapping with a default value
   * will be inserted.
   * @param key
   * @return AvlTreeMap<Key, Value>::Value &
   */
  Value &getOrDefault(K key) {
    if (this->root == nullptr) {
      this->root = Node::from(key);
      this->count += 1;
      return this->root->value;
    } else {
      return this
          ->getNodeOrProvide(this->root, key,
                             [&key]() { return Node::from(key); })
          ->value;
    }
  }

  /**
   * Returns true if this map contains a mapping for the specified key.
   * @param key
   * @return bool
   */
  bool contains(K key) const {
    return this->getNode(this->root, key) != nullptr;
  }

  /**
   * Associates the specified value with the specified key in this map.
   * @param key
   * @param value
   */
  void insert(K key, V value) {
    if (this->root == nullptr) {
      this->root = Node::from(key, value);
      this->count += 1;
    } else {
      this->insert(this->root, key, value);
    }
  }

  /**
   * If the specified key is not already associated with a value, associates
   * it with the given value and returns true, else returns false.
   * @param key
   * @param value
   * @return bool
   */
  bool insertIfAbsent(K key, V value) {
    USize sizeBeforeInsertion = this->size();
    if (this->root == nullptr) {
      this->root = Node::from(key, value);
      this->count += 1;
    } else {
      this->insert(this->root, key, value, false);
    }
    return this->size() > sizeBeforeInsertion;
  }

  /**
   * If the specified key is not already associated with a value, associates
   * it with the given value and returns the value, else returns the associated
   * value.
   * @param key
   * @param value
   * @return
   */
  Value &getOrInsert(K key, V value) {
    if (this->root == nullptr) {
      this->root = Node::from(key, value);
      this->count += 1;
      return root->value;
    } else {
      NodePtr node = getNodeOrProvide(this->root, key,
                                      [&]() { return Node::from(key, value); });
      return node->value;
    }
  }

  Value operator[](K key) const { return this->get(key); }

  Value &operator[](K key) { return this->getOrDefault(key); }

  /**
   * Removes the mapping for a key from this map if it is present;
   * Returns true if the mapping is present else returns false
   * @param key the key of the mapping
   * @return bool
   */
  bool remove(K key) {
    if (this->root == nullptr) {
      return false;
    } else {
      return this->remove(this->root, key, [](ConstNodePtr) {});
    }
  }

  /**
   * Removes the mapping for a key from this map if it is present and returns
   * the value which is mapped to the key; If this map contains no mapping for
   * the key, a {@code NoSuchMappingException} will be thrown.
   * @param key
   * @return AvlTreeMap<Key, Value>::Value
   * @throws NoSuchMappingException
   */
  Value getAndRemove(K key) {
    Value result;
    NodeConsumer action = [&](ConstNodePtr node) { result = node->value; };

    if (root == nullptr) {
      throw NoSuchMappingException("Invalid key");
    } else {
      if (remove(this->root, key, action)) {
        return result;
      } else {
        throw NoSuchMappingException("Invalid key");
      }
    }
  }

  /**
   * Gets the entry corresponding to the specified key; if no such entry
   * exists, returns the entry for the least key greater than the specified
   * key; if no such entry exists (i.e., the greatest key in the Tree is less
   * than the specified key), a {@code NoSuchMappingException} will be thrown.
   * @param key
   * @return AvlTreeMap<Key, Value>::Entry
   * @throws NoSuchMappingException
   */
  Entry getCeilingEntry(K key) const {
    if (this->root == nullptr) {
      throw NoSuchMappingException("No ceiling entry in this map");
    }

    NodePtr node = this->root;
    std::stack<NodePtr> ancestors;

    while (node != nullptr) {
      if (key == node->key) {
        return node->entry();
      }

      if (compare(key, node->key)) {
        /* key < node->key */
        if (node->left != nullptr) {
          ancestors.push(node);
          node = node->left;
        } else {
          return node->entry();
        }
      } else {
        /* key > node->key */
        if (node->right != nullptr) {
          ancestors.push(node);
          node = node->right;
        } else {
          if (ancestors.empty()) {
            throw NoSuchMappingException("No ceiling entry in this map");
          }

          NodePtr parent = ancestors.top();
          ancestors.pop();

          while (node == parent->right) {
            node = parent;
            if (!ancestors.empty()) {
              parent = ancestors.top();
              ancestors.pop();
            } else {
              throw NoSuchMappingException("No ceiling entry in this map");
            }
          }

          return parent->entry();
        }
      }
    }

    throw NoSuchMappingException("No ceiling entry in this map");
  }

  /**
   * Gets the entry corresponding to the specified key; if no such entry exists,
   * returns the entry for the greatest key less than the specified key;
   * if no such entry exists, a {@code NoSuchMappingException} will be thrown.
   * @param key
   * @return AvlTreeMap<Key, Value>::Entry
   * @throws NoSuchMappingException
   */
  Entry getFloorEntry(K key) const {
    if (this->root == nullptr) {
      throw NoSuchMappingException("No floor entry exists in this map");
    }

    NodePtr node = this->root;
    std::stack<NodePtr> ancestors;

    while (node != nullptr) {
      if (key == node->key) {
        return node->entry();
      }

      if (compare(key, node->key)) {
        /* key < node->key */
        if (node->left != nullptr) {
          ancestors.push(node);
          node = node->left;
        } else {
          if (ancestors.empty()) {
            throw NoSuchMappingException("No floor entry exists in this map");
          }

          NodePtr parent = ancestors.top();
          ancestors.pop();

          while (node == parent->left) {
            node = parent;
            if (!ancestors.empty()) {
              parent = ancestors.top();
              ancestors.pop();
            } else {
              throw NoSuchMappingException("No floor entry exists in this map");
            }
          }

          return parent->entry();
        }
      } else {
        /* key > node->key */
        if (node->right != nullptr) {
          ancestors.push(node);
          node = node->right;
        } else {
          return node->entry();
        }
      }
    }

    throw NoSuchMappingException("No floor entry exists in this map");
  }

  /**
   * Gets the entry for the least key greater than the specified
   * key; if no such entry exists, returns the entry for the least
   * key greater than the specified key; if no such entry exists,
   * a {@code NoSuchMappingException} will be thrown.
   * @param key
   * @return AvlTreeMap<Key, Value>::Entry
   * @throws NoSuchMappingException
   */
  Entry getHigherEntry(K key) {
    if (this->root == nullptr) {
      throw NoSuchMappingException("No higher entry exists in this map");
    }

    NodePtr node = this->root;
    std::stack<NodePtr> ancestors;

    while (node != nullptr) {
      if (compare(key, node->key)) {
        /* key < node->key */
        if (node->left != nullptr) {
          ancestors.push(node);
          node = node->left;
        } else {
          return node->entry();
        }
      } else {
        /* key >= node->key */
        if (node->right != nullptr) {
          ancestors.push(node);
          node = node->right;
        } else {
          if (ancestors.empty()) {
            throw NoSuchMappingException("No higher entry exists in this map");
          }

          NodePtr parent = ancestors.top();
          ancestors.pop();

          while (node == parent->right) {
            node = parent;
            if (!ancestors.empty()) {
              parent = ancestors.top();
              ancestors.pop();
            } else {
              throw NoSuchMappingException(
                  "No higher entry exists in this map");
            }
          }

          return parent->entry();
        }
      }
    }

    throw NoSuchMappingException("No higher entry exists in this map");
  }

  /**
   * Returns the entry for the greatest key less than the specified key; if
   * no such entry exists (i.e., the least key in the Tree is greater than
   * the specified key), a {@code NoSuchMappingException} will be thrown.
   * @param key
   * @return AvlTreeMap<Key, Value>::Entry
   * @throws NoSuchMappingException
   */
  Entry getLowerEntry(K key) const {
    if (this->root == nullptr) {
      throw NoSuchMappingException("No lower entry exists in this map");
    }

    NodePtr node = this->root;
    std::stack<NodePtr> ancestors;

    while (node != nullptr) {
      if (compare(key, node->key) || key == node->key) {
        /* key <= node->key */
        if (node->left != nullptr) {
          ancestors.push(node);
          node = node->left;
        } else {
          if (ancestors.empty()) {
            throw NoSuchMappingException("No lower entry exists in this map");
          }

          NodePtr parent = ancestors.top();
          ancestors.pop();

          while (node == parent->left) {
            node = parent;
            if (!ancestors.empty()) {
              parent = ancestors.top();
              ancestors.pop();
            } else {
              throw NoSuchMappingException("No lower entry exists in this map");
            }
          }

          return parent->entry();
        }
      } else {
        /* key > node->key */
        if (node->right != nullptr) {
          ancestors.push(node);
          node = node->right;
        } else {
          return node->entry();
        }
      }
    }

    throw NoSuchMappingException("No lower entry exists in this map");
  }

  /**
   * Remove all entries that satisfy the filter condition.
   * @param filter
   */
  void removeAll(KeyValueFilter filter) {
    std::vector<Key> keys;
    this->inorderTraversal([&](ConstNodePtr node) {
      if (filter(node->key, node->value)) {
        keys.push_back(node->key);
      }
    });
    for (const Key &key : keys) {
      this->remove(key);
    }
  }

  /**
   * Performs the given action for each key and value entry in this map.
   * The value is immutable for the action.
   * @param action
   */
  void forEach(KeyValueConsumer action) const {
    this->inorderTraversal(
        [&](ConstNodePtr node) { action(node->key, node->value); });
  }

  /**
   * Performs the given action for each key and value entry in this map.
   * The value is mutable for the action.
   * @param action
   */
  void forEachMut(MutKeyValueConsumer action) {
    this->inorderTraversal(
        [&](ConstNodePtr node) { action(node->key, node->value); });
  }

  /**
   * Returns a list containing all of the entries in this map.
   * @return AvlTreeMap<Key, Value>::EntryList
   */
  EntryList toEntryList() const {
    EntryList entryList;
    this->inorderTraversal(
        [&](ConstNodePtr node) { entryList.push_back(node->entry()); });
    return entryList;
  }

 private:
  static NodePtr rotateLeft(ConstNodePtr node) {
    // clang-format off
    //     |                       |
    //     N                       S
    //    / \     l-rotate(N)     / \
    //   L   S    ==========>    N   R
    //      / \                 / \
    //     M   R               L   M
    NodePtr successor = node->right;
    // clang-format on
    node->right = successor->left;
    successor->left = node;

    node->updateHeight();
    successor->updateHeight();

    return successor;
  }

  static NodePtr rotateRight(ConstNodePtr node) {
    // clang-format off
    //       |                   |
    //       N                   S
    //      / \   r-rotate(N)   / \
    //     S   R  ==========>  L   N
    //    / \                     / \
    //   L   M                   M   R
    NodePtr successor = node->left;
    // clang-format on
    node->left = successor->right;
    successor->right = node;

    node->updateHeight();
    successor->updateHeight();

    return successor;
  }

  static void swapNode(NodePtr &lhs, NodePtr &rhs) {
    std::swap(lhs->key, rhs->key);
    std::swap(lhs->value, rhs->value);
    std::swap(lhs, rhs);
  }

  static void fixBalance(NodePtr &node) {
    if (node->factor() < -1) {
      if (node->left->factor() < 0) {
        // clang-format off
        //  Left-Left Case
        //       |
        //       C                 |
        //      /   r-rotate(C)    B
        //     B    ==========>   / \
        //    /                  A   C
        //   A
        // clang-format on
        node = rotateRight(node);
      } else {
        // clang-format off
        //  Left-Right Case
        //     |                   |
        //     C                   C                 |
        //    /   l-rotate(A)     /   r-rotate(C)    B
        //   A    ==========>    B    ==========>   / \
        //    \                 /                  A   C
        //     B               A
        // clang-format on
        node->left = rotateLeft(node->left);
        node = rotateRight(node);
      }
    } else if (node->factor() > 1) {
      if (node->right->factor() > 0) {
        // clang-format off
        //  Right-Right Case
        //   |
        //   C                     |
        //    \     l-rotate(C)    B
        //     B    ==========>   / \
        //      \                A   C
        //       A
        // clang-format on
        node = rotateLeft(node);
      } else {
        // clang-format off
        //  Right-Left Case
        //   |                 |
        //   A                 A                     |
        //    \   r-rotate(C)   \     l-rotate(A)    B
        //     C  ==========>    B    ==========>   / \
        //    /                   \                A   C
        //   B                     C
        // clang-format on
        node->right = rotateRight(node->right);
        node = rotateLeft(node);
      }
    }
  }

  NodePtr getNodeOrProvide(NodePtr &node, K key, NodeProvider provide) {
    assert(node != nullptr);

    if (key == node->key) {
      return node;
    }

    assert(key != node->key);

    NodePtr result;

    if (compare(key, node->key)) {
      /* key < node->key */
      if (node->left == nullptr) {
        result = node->left = provide();
        this->count += 1;
        node->updateHeight();
      } else {
        result = getNodeOrProvide(node->left, key, provide);
        node->updateHeight();
        fixBalance(node);
      }
    } else {
      /* key > node->key */
      if (node->right == nullptr) {
        result = node->right = provide();
        this->count += 1;
        node->updateHeight();
      } else {
        result = getNodeOrProvide(node->right, key, provide);
        node->updateHeight();
        fixBalance(node);
      }
    }

    return result;
  }

  NodePtr getNode(ConstNodePtr node, K key) const {
    assert(node != nullptr);

    if (key == node->key) {
      return node;
    }

    if (compare(key, node->key)) {
      /* key < node->key */
      return node->left == nullptr ? nullptr : getNode(node->left, key);
    } else {
      /* key > node->key */
      return node->right == nullptr ? nullptr : getNode(node->right, key);
    }
  }

  void insert(NodePtr &node, K key, V value, bool replace = true) {
    assert(node != nullptr);

    if (key == node->key) {
      if (replace) {
        node->value = value;
      }
      return;
    }

    assert(key != node->key);

    if (compare(key, node->key)) {
      /* key < node->key */
      if (node->left == nullptr) {
        node->left = Node::from(key, value);
        this->count += 1;
        node->updateHeight();
      } else {
        insert(node->left, key, value, replace);
        node->updateHeight();
        fixBalance(node);
      }
    } else {
      /* key > node->key */
      if (node->right == nullptr) {
        node->right = Node::from(key, value);
        this->count += 1;
        node->updateHeight();
      } else {
        insert(node->right, key, value, replace);
        node->updateHeight();
        fixBalance(node);
      }
    }
  }

  bool remove(NodePtr &node, K key, NodeConsumer action) {
    assert(node != nullptr);

    if (key != node->key) {
      if (compare(key, node->key)) {
        /* key < node->key */
        NodePtr &left = node->left;
        if (left != nullptr && remove(left, key, action)) {
          node->updateHeight();
          fixBalance(node);
          return true;
        } else {
          return false;
        }
      } else {
        /* key > node->key */
        NodePtr &right = node->right;
        if (right != nullptr && remove(right, key, action)) {
          node->updateHeight();
          fixBalance(node);
          return true;
        } else {
          return false;
        }
      }
    }

    assert(key == node->key);
    action(node);

    if (node->isLeaf()) {
      // Case 1: no child
      node = nullptr;
    } else if (node->right == nullptr) {
      // clang-format off
      // Case 2: left child only
      //     P
      //     |  remove(N)  P
      //     N  ========>  |
      //    /              L
      //   L
      // clang-format on
      node = node->left;
      node->updateHeight();
    } else if (node->left == nullptr) {
      // clang-format off
      // Case 3: right child only
      //   P
      //   |    remove(N)  P
      //   N    ========>  |
      //    \              R
      //     R
      // clang-format on
      node = node->right;
      node->updateHeight();
    } else if (node->right->left == nullptr) {
      // clang-format off
      // Case 4: both left and right child, right child has no left child
      //    |                 |
      //    N    remove(N)    R
      //   / \   ========>   /
      //  L   R             L
      // clang-format on
      NodePtr right = node->right;
      swapNode(node, right);
      right->right = node->right;
      node = right;
      node->updateHeight();
      fixBalance(node);
    } else {
      // clang-format off
      // Case 5: both left and right child, right child is not a leaf
      //   Step 1. find the node N with the smallest key
      //           and its parent P on the right subtree
      //   Step 2. swap S and N
      //   Step 3. remove node N like Case 1 or Case 3
      //   Step 4. update height for P
      //     |                  |
      //     N                  S                 |
      //    / \                / \                S
      //   L  ..  swap(N, S)  L  ..  remove(N)   / \
      //       |  =========>      |  ========>  L  ..
      //       P                  P                 |
      //      / \                / \                P
      //     S  ..              N  ..              / \
      //      \                  \                R  ..
      //       R                  R
      // clang-format on

      // Step 1
      NodePtr successor = node->right;
      NodePtr parent = node;
      while (successor->left != nullptr) {
        parent = successor;
        successor = parent->left;
      }
      // Step 2
      swapNode(node, successor);
      // Step 3
      parent->left = node->right;
      // Restore node
      node = successor;
      // Step 4
      parent->updateHeight();
    }

    this->count -= 1;
    return true;
  }

  void inorderTraversal(NodeConsumer action) const {
    if (this->root == nullptr) {
      return;
    }

    std::stack<NodePtr> stack;
    NodePtr node = this->root;

    while (node != nullptr || !stack.empty()) {
      while (node != nullptr) {
        stack.push(node);
        node = node->left;
      }
      if (!stack.empty()) {
        node = stack.top();
        stack.pop();
        action(node);
        node = node->right;
      }
    }
  }
};

#endif  // AVLTREE_MAP_HPP

Tài liệu khác

Tại AVL Tree Visualization có thể quan sát quá trình duy trì sự cân bằng của cây AVL.

Wikipedia -- AVL Tree

---

Last updated on this page:, Update history
Found an error? Want to help improve? Edit this page on GitHub!
Contributors to this page:OI-wiki
All content on this page is provided under the terms of the CC BY-SA 4.0 and SATA license, additional terms may apply