AVL树C++代码实现

前言

C++实现的一个平衡二叉树， 参考了许夜3大佬的这篇文章，建议读完后再看我的代码。

这个类只是我项目的一部分，有些成员的实现不在这里，不用担心，不影响阅读，看到跳过就行。

如果对我的项目感兴趣的话，可以看这个链接，如果你觉得我写的不错，麻烦给俺个star~

代码实现

template<typename Key, typename Val, typename Compare = Less<Key>, typename Equal_ = Equal<Key>>
    class AVL_Tree {
    public:
        friend class Detail::AVL_Tree_Iterator<Key, Val, Compare, Equal_>;

        using Iterator = Detail::AVL_Tree_Iterator<Key, Val, Compare, Equal_>;

        using cIterator = STD::cIterator<Iterator, Bidirectional_iterator_tag>;

        using rIterator = STD::rIterator<Iterator, Bidirectional_iterator_tag>;

        using crIterator = STD::crIterator<Iterator, Bidirectional_iterator_tag>;

        using Self = AVL_Tree<Key, Val, Compare, Equal_>;

    private:
        struct Node {
            Pair<Key, Val> value;

            Node *left = nullptr, *right = nullptr, *last = nullptr;

            int height = 1;

            explicit Node(const Key &key, const Val &val, int height = 1) :
                    value({key, val}), height(height) {};

            explicit Node(const Key &key, Val &&val, int height = 0) :
                    value({key, move(val)}), height(height) {};

            const Key &key() const { return value.first; };

        };

        Node *root = nullptr, *val_begin = nullptr, *val_end = nullptr;

        Compare less = Less<Key>();

        Equal_ equal = Equal<Key>();

        Size size_ = 0;

        // 计算传入节点的左节点与右节点之差
        int height_difference(Node *node) const {
            if (!node) return 0;
            int left = node->left ? node->left->height : 0;
            int right = node->right ? node->right->height : 0;
            return left - right;
        };

        // 重新计算传入节点的高度
        void calculate_height(Node *node) const {
            int left = node->left ? node->left->height : 0;
            int right = node->right ? node->right->height : 0;
            node->height = left > right ? left + 1 : right + 1;
        }

        // link操作是为了更新子节点的last，并将子节点连接到父节点上，同时会重新计算父节点的高度，可以减少许多麻烦。
        void link(Node *parent, bool left, Node *child) const {
            if (parent) {
                if (left) parent->left = child;
                else parent->right = child;
                calculate_height(parent);
            }
            if (child) child->last = parent;
        };

        // 左旋
        Node *L_alpha(Node *node) {
            Node *n_right = node->right;
            link(node, false, n_right->left);
            link(n_right, true, node);
            return n_right;
        }

        // 右旋
        Node *R_alpha(Node *node) {
            Node *n_left = node->left;
            link(node, true, n_left->right);
            link(n_left, false, node);
            return n_left;
        }

        // 该函数的作用是在树被破坏时，重新平衡树。
        Node *balance(Node *node);

        // 内置版本的插入函数。
        Node *insert(Node *node, Node *&target);

        // 找到该节点下的最小值。
        static Node *the_min(Node *node);

        // 找到该节点下的最大值。
        static Node *the_max(Node *node);

        // 删除该节点下的最大值。
        Node *erase_max(Node *node);

        // 内置递归版本的删除函数。
        Node *erase(Node *node, const Key &key);

        // 将该节点的全部分支删除。
        void erase_from(Node *node);

    public:
        AVL_Tree() = default;

        ~AVL_Tree();

        void clear();

        Pair<Iterator, bool> insert(const Key &key, const Val &value);

        Pair<Iterator, bool> insert(const Key &key, Val &&value);

        bool erase(const Key &key);

        Iterator find(const Key &key) const;

    };
    
//----------------------------------------------------------------------------------------------------------------------

    template<typename Key, typename Val, typename Compare, typename Equal>
    typename AVL_Tree<Key, Val, Compare, Equal>::Node *
    AVL_Tree<Key, Val, Compare, Equal>::the_min(Node *node) {
        if (!node) return nullptr;
        while (node->left) node = node->left;
        return node;
    };

    template<typename Key, typename Val, typename Compare, typename Equal>
    typename AVL_Tree<Key, Val, Compare, Equal>::Node *
    AVL_Tree<Key, Val, Compare, Equal>::the_max(Node *node) {
        if (!node) return nullptr;
        while (node->right) node = node->right;
        return node;
    };

    template<typename Key, typename Val, typename Compare, typename Equal_>
    void AVL_Tree<Key, Val, Compare, Equal_>::erase_from(Node *node) {
        if (!node) return;
        erase_from(node->left);
        erase_from(node->right);
        Deallocate(node);
        --size_;
    };

    template<typename Key, typename Val, typename Compare, typename Equal_>
    typename AVL_Tree<Key, Val, Compare, Equal_>::Node
    *AVL_Tree<Key, Val, Compare, Equal_>::balance(Node *node) {
        int difference = height_difference(node); // 得到左右的节点高度差。
        //一共有六种情况（插入会出现四种，删除会出现六种），这里合并为四种。
        if (difference > 1) { // 本节点的左边太长
            if (height_difference(node->left) >= 0) {
                node = R_alpha(node);
            } else { // 本节点的左节点的右边比较长
                link(node, true, L_alpha(node->left));
                node = R_alpha(node);
            }
        } else if (difference < -1) { // 本节点的右边太长
            if (height_difference(node->right) <= 0) {
                node = L_alpha(node);
            } else { // 本节点的右节点的左边比较长
                link(node, false, R_alpha(node->right));
                node = L_alpha(node);
            }
        }
        return node;
    };

    template<typename Key, typename Val, typename Compare, typename Equal_>
    typename AVL_Tree<Key, Val, Compare, Equal_>::Node
    *AVL_Tree<Key, Val, Compare, Equal_>::insert(Node *node, Node *&target) {
        if (!node) { // 说明找到了合适的位置插入，将目标向上传递，终止递归。
            ++size_;
            return target;
        }
        if (less(target->key(), node->key())) // 目标小于本节点，向左边寻找。
            link(node, true, insert(node->left, target));
        else if (!equal(target->key(), node->key())) // 目标大于本节点，向右边寻找。
            link(node, false, insert(node->right, target));
        else { // 找到了相等值，给本节点重新赋值。
            node->value = move(target->value);
            Deallocate(target);
            target = node;
        }
        return balance(node);
    };

    template<typename Key, typename Val, typename Compare, typename Equal_>
    typename AVL_Tree<Key, Val, Compare, Equal_>::Node
    *AVL_Tree<Key, Val, Compare, Equal_>::erase_max(Node *node) {
        if (node->right) { // 当没有找到最大值时，继续递归寻找。
            link(node, false, erase_max(node->right));
        } else if (node->left) { // 找到最大值，但其还有左节点，我们将其值换上来，删掉左节点。
            STD::swap(node->value, node->left->value);
            link(node, true, erase_max(node->left));
        } else { // 找到最大值，并且它没有任何子节点，可直接删除。
            if (val_end == node) val_end = node->last;
            Deallocate(node);
            --size_;
            return nullptr;
        }
        return balance(node);
    };

    template<typename Key, typename Val, typename Compare, typename Equal_>
    typename AVL_Tree<Key, Val, Compare, Equal_>::Node
    *AVL_Tree<Key, Val, Compare, Equal_>::erase(Node *node, const Key &key) {
        if (!node) return nullptr; // 说明树中没有该目标，终止递归。
        if (less(key, node->key())) // 目标值小于当前值，向左继续寻找。
            link(node, true, erase(node->left, key));
        else if (!equal(key, node->key())) // 目标值大于当前值，向右继续寻找。
            link(node, false, erase(node->right, key));
        else { // 找到目标值
            if (node->left) { // 有子节点比较麻烦，我们选择置换上来当前节点左边的最大值，转而删除被置换的那个节点。
                STD::swap(node->value, the_max(node->left)->value);
                link(node, true, erase_max(node->left));
            } else { // 这里合并了两种情况（有右节点和无节点），直接删除该节点，将该节点的右节点向上传递。
                auto temp = node->right;
                if (node == val_begin) val_begin = temp;
                if (node == val_end) val_end = node->last;
                Deallocate(node);
                --size_;
                return temp;
            }
        }
        return balance(node);
    };

//----------------------------------------------------------------------------------------------------------------------

    template<typename Key, typename Val, typename Compare, typename Equal_>
    AVL_Tree<Key, Val, Compare, Equal_>::~AVL_Tree() {
        erase_from(root);
    }

    template<typename Key, typename Val, typename Compare, typename Equal_>
    Pair<typename AVL_Tree<Key, Val, Compare, Equal_>::Iterator, bool>
    AVL_Tree<Key, Val, Compare, Equal_>::insert(const Key &key, const Val &value) {
        Size record = size_;
        // 节点默认高度为1。
        Node *temp = Allocate<Node>(key, value, 1);
        if (!size_) {
            val_begin = temp;
            val_end = temp;
        } else if (less(key, val_begin->key())) {
            val_begin = temp;
        } else if (!less(key, val_end->key()) && !equal(key, val_end->key())) {
            val_end = temp;
        }
        root = insert(root, temp);
        root->last = nullptr;
        return {Iterator(temp, this), size_ != record};
    };

    template<typename Key, typename Val, typename Compare, typename Equal_>
    bool AVL_Tree<Key, Val, Compare, Equal_>::erase(const Key &key) {
        Size record = size_;
        root = erase(root, key);
        if (root) root->last = nullptr;
        return record != size_;
    };