AVLTree/Linux/AVLTree.hpp

#pragma once

#include <iostream>
#include <stack>
#include <queue>

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

        int _bf;             // balance factor	平衡因子
        std::pair<K, V> _kv; // key-value

        // 普通构造
        AVLTreeNode(const std::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 std::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);
        //	std::cout << root->_kv.first << ":" << root->_kv.second << std::endl;
        //	_InOrder(root->_right);
        //}

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

        // 中序遍历 非递归
        void InOrder()
        {
            std::stack<Node *> stack;
            Node *cur = _root;

            while (cur != nullptr || !stack.empty())
            {
                // 找到最左边的节点
                while (cur != nullptr)
                {
                    stack.push(cur);
                    cur = cur->_left;
                }

                // 处理当前节点
                cur = stack.top();
                stack.pop();
                std::cout << cur->_kv.first << ":" << cur->_kv.second << std::endl;

                // 转向右子树
                cur = cur->_right;
            }
            std::cout << std::endl;
        }

        // 前序遍历的递归实现
        // 辅助函数
        // void _PreOrder(Node* root)
        //{
        //	if (root == nullptr)
        //	{
        //		return;
        //	}

        //	std::cout << root->_kv.first << ":" << root->_kv.second << std::endl;   // 访问根节点
        //	_PreOrder(root->_left);   // 访问左子树
        //	_PreOrder(root->_right);  // 访问右子树
        //}

        // void PreOrder()
        //{
        //	_PreOrder(_root);
        //	std::cout << std::endl;
        // }

        // 前序遍历的非递归实现
        void PreOrder()
        {
            if (_root == nullptr)
            {
                return;
            }

            std::stack<Node *> stack;
            stack.push(_root);

            while (!stack.empty())
            {
                Node *cur = stack.top();
                stack.pop();
                std::cout << cur->_kv.first << ":" << cur->_kv.second << std::endl; // 访问当前节点

                // 先压入右子树再压入左子树（确保左子树先处理）
                if (cur->_right != nullptr)
                {
                    stack.push(cur->_right);
                }
                if (cur->_left != nullptr)
                {
                    stack.push(cur->_left);
                }
            }
            std::cout << std::endl;
        }

        // 后序遍历的递归实现
        // 辅助函数
        // void _PostOrder(Node* root) {
        //	if (root == nullptr)
        //	{
        //		return;
        //	}

        //	_PostOrder(root->_left);    // 访问左子树
        //	_PostOrder(root->_right);   // 访问右子树
        //	std::cout << root->_kv.first << ":" << root->_kv.second << std::endl;     // 访问根节点
        //}

        // void PostOrder()
        //{
        //	_PostOrder(_root);
        //	std::cout << std::endl;
        // }

        // 后序遍历的非递归实现
        void PostOrder()
        {
            if (_root == nullptr)
            {
                return;
            }

            std::stack<Node *> st;
            Node *cur = _root;
            Node *lastNode = nullptr; // 最近访问的节点
            while (cur || !st.empty())
            {
                while (cur)
                {
                    st.push(cur);
                    cur = cur->_left;
                }

                Node *top = st.top();
                if ((top->_right == nullptr) || (lastNode == top->_right))
                {
                    std::cout << top->_kv.first << ":" << top->_kv.second << std::endl;
                    lastNode = top;
                    st.pop();
                }
                else
                {
                    cur = top->_right;
                }
            }
            std::cout << std::endl;
        }

        // 层序遍历
        void LevelOrder()
        {
            if (_root == nullptr)
            {
                return;
            }

            std::queue<Node *> queue;
            queue.push(_root);

            while (!queue.empty())
            {
                Node *cur = queue.front();
                queue.pop();
                std::cout << cur->_kv.first << ":" << cur->_kv.second << std::endl; // 访问当前节点

                if (cur->_left != nullptr)
                {
                    queue.push(cur->_left); // 将左子树压入队列
                }
                if (cur->_right != nullptr)
                {
                    queue.push(cur->_right); // 将右子树压入队列
                }
            }
            std::cout << std::endl;
        }

        // 查找后继节点
        Node *findMin(Node *node)
        {
            while (node->_left != nullptr)
            {
                node = node->_left;
            }
            return node;
        }

        Node *FindSuccessor(Node *node)
        {
            if (node->_right != nullptr)
            {
                return findMin(node->_right);
            }

            Node *cur = _root;
            Node *successor = nullptr;
            while (cur != nullptr)
            {
                if (node->_kv.first < cur->_kv.first)
                {
                    successor = cur;
                    cur = cur->_left;
                }
                else if (node->_kv.first > cur->_kv.first)
                {
                    cur = cur->_right;
                }
                else
                {
                    break;
                }
            }
            return successor;
        }

        // 查找前驱节点
        Node *findMax(Node *node)
        {
            while (node->_right != nullptr)
            {
                node = node->_right;
            }
            return node;
        }

        Node *FindPredecessor(Node *node)
        {
            if (node->_left != nullptr)
            {
                return findMax(node->_left);
            }

            Node *cur = _root;
            Node *predecessor = nullptr;
            while (cur != nullptr)
            {
                if (node->_kv.first < cur->_kv.first)
                {
                    cur = cur->_left;
                }
                else if (node->_kv.first > cur->_kv.first)
                {
                    predecessor = cur;
                    cur = cur->_right;
                }
                else
                {
                    break;
                }
            }
            return predecessor;
        }

        // 获取树的高度
        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 K& key)
        //{
        //	Node* cur = _root;
        //	while (cur)
        //	{
        //		if (cur->_kv.first > key)
        //		{
        //			cur = cur->_left;
        //		}
        //		else if (cur->_kv.first < key)
        //		{
        //			cur = cur->_right;
        //		}
        //		else
        //		{
        //			return cur;
        //		}
        //	}
        //	return nullptr;
        //}

        // 查找的递归实现
        // 辅助函数
        Node *_find(Node *root, const K &key)
        {
            if (root == nullptr || root->_kv.first == key)
            {
                return root;
            }
            if (key < root->_kv.first)
            {
                return _find(root->_left, key);
            }
            else
            {
                return _find(root->_right, key);
            }
        }

        // 查找节点
        Node *Find(const K &key)
        {
            return _find(_root, key);
        }

    private:
        Node *_root = nullptr;
    };
}