RBTree/Linux/AVLTree.hpp

#pragma once
#include <iostream>

using namespace std;

namespace Lenyiin
{
    template <class K, class V>
    struct AVLTreeNode
    {
        AVLTreeNode<K, V> *_left;
        AVLTreeNode<K, V> *_right;
        AVLTreeNode<K, V> *_parent;

        int _bf;
        pair<K, V> _kv;

        // 普通构造
        AVLTreeNode(const pair<K, V> &kv)
            : _left(nullptr), _right(nullptr), _parent(nullptr), _bf(0), _kv(kv)
        {
        }
    };

    template <class K, class V>
    class AVLTree
    {
    private:
        typedef AVLTreeNode<K, V> Node;

    public:
        bool Insert(const pair<K, V> &kv)
        {
            // 1. 先按搜索树的规则进行插入
            if (_root == nullptr)
            {
                _root = new Node(kv);
                return true;
            }

            Node *parent = nullptr;
            Node *cur = _root;
            while (cur)
            {
                if (cur->_kv.first > kv.first)
                {
                    parent = cur;
                    cur = cur->_left;
                }
                else if (cur->_kv.first < kv.first)
                {
                    parent = cur;
                    cur = cur->_right;
                }
                else
                {
                    return false;
                }
            }

            // 找到位置
            cur = new Node(kv);
            if (parent->_kv.first > kv.first)
            {
                parent->_left = cur;
                cur->_parent = parent;
            }
            else
            {
                parent->_right = cur;
                cur->_parent = parent;
            }

            // 2. 更新平衡因子
            while (parent)
            {
                if (cur == parent->_right)
                {
                    parent->_bf++; // 插在右边，平衡因子++
                }
                else
                {
                    parent->_bf--; // 插在左边，平衡因子--
                }

                if (parent->_bf == 0)
                {
                    // 说明 parent 所在的子树的高度不变，更新结束
                    break;
                }
                else if (parent->_bf == 1 || parent->_bf == -1)
                {
                    // 说明 parent 所在的子树的高度变了，继续往上更新
                    cur = parent;
                    parent = parent->_parent;
                }
                else if (parent->_bf == 2 || parent->_bf == -2)
                {
                    // parent 所在的子树出现问题了，需要旋转处理一下
                    // 1. 旋转前提是保持它依旧是搜索二叉树
                    // 2. 旋转成平衡树
                    if (parent->_bf == 2)
                    {
                        if (cur->_bf == 1)
                        {
                            // 左旋
                            RotateL(parent);
                        }
                        else if (cur->_bf == -1)
                        {
                            // 右左双旋
                            RotateRL(parent);
                        }
                    }
                    else if (parent->_bf == -2)
                    {
                        if (cur->_bf == -1)
                        {
                            // 右旋
                            RotateR(parent);
                        }
                        else if (cur->_bf == 1)
                        {
                            // 左右双旋
                            RotateLR(parent);
                        }
                    }

                    // 旋转完成后, parent 所在的树的高度恢复到了，插入节点前高度
                    // 如果是子树，对上一层没有影响，更新结束
                    break;
                }
            }

            return true;
        }

        // 左单旋
        void RotateL(Node *parent)
        {
            Node *ppNode = parent->_parent;
            Node *subR = parent->_right;
            Node *subRL = subR->_left;

            parent->_right = subRL;
            if (subRL)
            {
                subRL->_parent = parent;
            }
            subR->_left = parent;
            parent->_parent = subR;

            // 1. 原来 parent 是这棵树的根，现在 subR 是根
            if (parent == _root)
            {
                _root = subR;
                subR->_parent = nullptr;
            }
            // 2. parent 为根的树只是整棵树中的子树，改变链接关系，那么 subR 要顶替他的位置
            else
            {
                if (ppNode->_left == parent)
                {
                    ppNode->_left = subR;
                }
                else
                {
                    ppNode->_right = subR;
                }
                subR->_parent = ppNode;
            }

            parent->_bf = subR->_bf = 0;
        }

        // 右单旋
        void RotateR(Node *parent)
        {
            Node *ppNode = parent->_parent;
            Node *subL = parent->_left;
            Node *subLR = subL->_right;

            parent->_left = subLR;
            if (subLR)
            {
                subLR->_parent = parent;
            }
            subL->_right = parent;
            parent->_parent = subL;

            if (parent == _root)
            {
                _root = subL;
                subL->_parent = nullptr;
            }
            else
            {
                if (ppNode->_left == parent)
                {
                    ppNode->_left = subL;
                }
                else
                {
                    ppNode->_right = subL;
                }
                subL->_parent = ppNode;
            }
            parent->_bf = subL->_bf = 0;
        }

        // 右左双旋
        void RotateRL(Node *parent)
        {
            Node *subR = parent->_right;
            Node *subRL = subR->_left;
            int bf = subRL->_bf;

            RotateR(parent->_right);
            RotateL(parent);

            if (bf == -1)
            {
                parent->_bf = 0;
                subR->_bf = 1;
                subRL->_bf = 0;
            }
            else if (bf == 1)
            {
                parent->_bf = -1;
                subR->_bf = 0;
                subRL->_bf = 0;
            }
            else if (bf == 0)
            {
                parent->_bf = 0;
                subR->_bf = 0;
                subRL->_bf = 0;
            }
        }

        // 左右双旋
        void RotateLR(Node *parent)
        {
            Node *subL = parent->_left;
            Node *subLR = subL->_right;
            int bf = subLR->_bf;

            RotateL(subL);
            RotateR(parent);

            if (bf == 1)
            {
                parent->_bf = 0;
                subL->_bf = -1;
                subLR->_bf = 0;
            }
            else if (bf == -1)
            {
                parent->_bf = 1;
                subL->_bf = 0;
                subLR->_bf = 0;
            }
            else if (bf == 0)
            {
                parent->_bf = 0;
                subL->_bf = 0;
                subLR->_bf = 0;
            }
        }

        // 中序遍历
        void _InOrder(Node *root)
        {
            if (root == nullptr)
            {
                return;
            }

            _InOrder(root->_left);
            cout << root->_kv.first << ":" << root->_kv.second << endl;
            _InOrder(root->_right);
        }

        void InOrder()
        {
            _InOrder(_root);
            cout << endl;
        }

        // 获取树的高度
        int Height(Node *root)
        {
            if (root == nullptr)
            {
                return 0;
            }

            int leftHeight = Height(root->_left);
            int rightHeight = Height(root->_right);

            return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
        }

        // 判断是否是平衡二叉树
        bool _IsBalance(Node *root)
        {
            if (root == nullptr)
            {
                return true;
            }

            int leftHeight = Height(root->_left);
            int rightHeight = Height(root->_right);
            if (abs(leftHeight - rightHeight) > 1)
            {
                return false;
            }

            return _IsBalance(root->_left) && _IsBalance(root->_right);
        }

        bool IsBalance()
        {
            return _IsBalance(_root);
        }

        // 查找
        Node *Find(const Node *root)
        {
            Node *cur = _root;
            while (cur)
            {
                if (cur->_kv.first > root->_kv.first)
                {
                    cur = cur->_left;
                }
                else if (cur->_kv.first < root->_kv.first)
                {
                    cur = cur->_right;
                }
                else
                {
                    return cur;
                }
            }
            return nullptr;
        }

    private:
        Node *_root = nullptr;
    };
}