RBTree/Linux/RBTree.hpp

#pragma once

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

namespace Lenyiin
{
    enum Colour
    {
        RED,
        BLACK
    };

    template <class K, class V>
    struct RBTreeNode
    {
        RBTreeNode<K, V> *_left;
        RBTreeNode<K, V> *_right;
        RBTreeNode<K, V> *_parent;

        std::pair<K, V> _kv;
        Colour _colour;

        RBTreeNode(const std::pair<K, V> &kv)
            : _left(nullptr), _right(nullptr), _parent(nullptr), _kv(kv), _colour(RED)
        {
        }
    };

    template <class K, class V>
    class RBTree
    {
    private:
        typedef RBTreeNode<K, V> Node;

        // 左单旋
        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;
            }
        }

        // 右单旋
        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;
            }
        }

        // 删除整个红黑树
        void DeleteTree(Node *root)
        {
            if (root == nullptr)
            {
                return;
            }

            // 递归删除左子树
            DeleteTree(root->_left);
            // 递归删除右子树
            DeleteTree(root->_right);

            // 删除当前节点
            delete root;
        }

    public:
        RBTree(Node *root = nullptr)
            : _root(root)
        {
        }

        ~RBTree()
        {
            DeleteTree(_root);
        }

        // 1. 空树。插入结点做根，把他变黑。
        // 2. 插入红色节点，他的父亲是黑色的，结束。
        // 3. 插入红色节点，他的父亲是红色的，可以推断他的祖父存在且一定为黑色。关键看叔叔
        //	a. 如果叔叔存在且为红，把父亲和叔叔变黑，祖父变红，继续往上处理。
        //	b. 如果叔叔存在且为黑，或者不存在。旋转（单旋 or 双旋）+ 变色
        bool Insert(const std::pair<K, V> &kv)
        {
            // 1. 按照搜索树的规则进行插入
            // 如果树为空，直接将新节点设为根节点，并染成黑色
            if (_root == nullptr)
            {
                _root = new Node(kv);
                _root->_colour = BLACK; // 根节点是黑色的
                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;
            }

            // 插入调整
            InsertFixUp(parent, cur);

            return true;
        }

        void InsertFixUp(Node *parent, Node *cur)
        {
            // 调整结点颜色
            // 新结点给红的还是黑的?	红色
            // 1. 空树。插入结点做根, 把他变黑。
            // 2. 插入红色节点, 他的父亲是黑色的, 结束。
            // 3. 插入红色节点, 他的父亲是红色的, 可以推断他的祖父存在且一定为黑色。关键看叔叔
            //	a. 如果叔叔存在且为红, 把父亲和叔叔变黑, 祖父变红, 继续往上处理。
            //	b. 如果叔叔存在且为黑, 或者不存在。旋转（单旋 or 双旋）+ 变色
            while (parent && parent->_colour == RED)
            {
                // 关键看叔叔
                Node *grandfather = parent->_parent;
                if (grandfather->_left == parent)
                {
                    Node *uncle = grandfather->_right;

                    // 情况1: uncle 存在, 且为红
                    if (uncle && uncle->_colour == RED)
                    {
                        parent->_colour = uncle->_colour = BLACK;
                        grandfather->_colour = RED;

                        // 继续向上调整
                        cur = grandfather;
                        parent = cur->_parent;
                    }
                    // 情况 2 or 情况 3 : uncle 不存在 or uncle 存在且为黑
                    else
                    {
                        // 情况 3: 双旋 -> 变成单旋
                        if (cur == parent->_right)
                        {
                            RotateL(parent);
                            std::swap(parent, cur);
                        }

                        // 第二种情况 (ps: 有可能是第三种情况变过来的)
                        RotateR(grandfather);
                        grandfather->_colour = RED;
                        parent->_colour = BLACK;

                        break;
                    }
                }
                else
                {
                    Node *uncle = grandfather->_left;

                    // 情况1: uncle 存在, 且为红
                    if (uncle && uncle->_colour == RED)
                    {
                        parent->_colour = BLACK;
                        uncle->_colour = BLACK;
                        grandfather->_colour = RED;

                        // 继续向上调整
                        cur = grandfather;
                        parent = cur->_parent;
                    }
                    // 情况2 or 情况3: uncle 不存在 or uncle 存在且为黑
                    else
                    {
                        // 情况3: 双旋 -> 变为单旋
                        if (cur == parent->_left)
                        {
                            RotateR(parent);
                            std::swap(parent, cur);
                        }

                        // 第二种情况 (ps: 有可能是第三种情况变来的)
                        RotateL(grandfather);
                        grandfather->_colour = RED;
                        parent->_colour = BLACK;

                        break;
                    }
                }
            }

            // 保证根节点为黑色
            _root->_colour = BLACK;
        }

        // 删除操作
        bool Erase(const K &key)
        {
            // 查找节点
            Node *nodeToDelete = _root;
            while (nodeToDelete)
            {
                if (nodeToDelete->_kv.first > key)
                {
                    nodeToDelete = nodeToDelete->_left;
                }
                else if (nodeToDelete->_kv.first < key)
                {
                    nodeToDelete = nodeToDelete->_right;
                }
                else
                {
                    break; // 找到待删除的节点
                }
            }

            // 如果未找到节点
            if (nodeToDelete == nullptr)
            {
                return false;
            }

            // 执行删除操作
            Node *parent, *child;
            Colour originalColour = nodeToDelete->_colour;

            if (nodeToDelete->_left == nullptr)
            {
                child = nodeToDelete->_right;
                parent = nodeToDelete->_parent;
                Transplant(nodeToDelete, nodeToDelete->_right);
            }
            else if (nodeToDelete->_right == nullptr)
            {
                child = nodeToDelete->_left;
                parent = nodeToDelete->_parent;
                Transplant(nodeToDelete, nodeToDelete->_left);
            }
            else
            {
                Node *successor = Minimum(nodeToDelete->_right);
                originalColour = successor->_colour;
                child = successor->_right;
                if (successor->_parent == nodeToDelete)
                {
                    parent = successor;
                }
                else
                {
                    Transplant(successor, successor->_right);
                    successor->_right = nodeToDelete->_right;
                    successor->_right->_parent = successor;
                    parent = successor->_parent;
                }
                Transplant(nodeToDelete, successor);
                successor->_left = nodeToDelete->_left;
                successor->_left->_parent = successor;
                successor->_colour = nodeToDelete->_colour;
            }

            delete nodeToDelete;

            // 如果删除的节点是黑色，需要进行调整
            if (originalColour == BLACK)
            {
                EraseFixUp(child, parent);
            }

            return true;
        }

        void EraseFixUp(Node *x, Node *parent)
        {
            while (x != _root && (x == nullptr || x->_colour == BLACK))
            {
                if (x == parent->_left)
                {
                    Node *w = parent->_right;
                    // 情况1: x的兄弟节点w是红色的
                    if (w->_colour == RED)
                    {
                        w->_colour = BLACK;
                        parent->_colour = RED;
                        RotateL(parent);
                        w = parent->_right;
                    }
                    // 情况2: x的兄弟节点w是黑色的，且w的两个子节点都是黑色的
                    if ((w->_left == nullptr || w->_left->_colour == BLACK) &&
                        (w->_right == nullptr || w->_right->_colour == BLACK))
                    {
                        w->_colour = RED;
                        x = parent;
                        parent = x->_parent;
                    }
                    else
                    {
                        // 情况3: w是黑色的，并且w的左子节点是红色，右子节点是黑色
                        if (w->_right == nullptr || w->_right->_colour == BLACK)
                        {
                            if (w->_left)
                                w->_left->_colour = BLACK;
                            w->_colour = RED;
                            RotateR(w);
                            w = parent->_right;
                        }
                        // 情况4: w是黑色的，并且w的右子节点是红色
                        if (w)
                        {
                            w->_colour = parent->_colour;
                            parent->_colour = BLACK;
                            if (w->_right)
                                w->_right->_colour = BLACK;
                            RotateL(parent);
                        }
                        x = _root;
                        break;
                    }
                }
                else
                {
                    Node *w = parent->_left;
                    // 情况1: x的兄弟节点w是红色的
                    if (w->_colour == RED)
                    {
                        w->_colour = BLACK;
                        parent->_colour = RED;
                        RotateR(parent);
                        w = parent->_left;
                    }
                    // 情况2: x的兄弟节点w是黑色的，且w的两个子节点都是黑色的
                    if ((w->_left == nullptr || w->_left->_colour == BLACK) &&
                        (w->_right == nullptr || w->_right->_colour == BLACK))
                    {
                        w->_colour = RED;
                        x = parent;
                        parent = x->_parent;
                    }
                    else
                    {
                        // 情况3: w是黑色的，并且w的右子节点是红色，左子节点是黑色
                        if (w->_left == nullptr || w->_left->_colour == BLACK)
                        {
                            if (w->_right)
                            {
                                w->_right->_colour = BLACK;
                            }
                            w->_colour = RED;
                            RotateL(w);
                            w = parent->_left;
                        }
                        // 情况4: w是黑色的，并且w的左子节点是红色
                        if (w)
                        {
                            w->_colour = parent->_colour;
                            parent->_colour = BLACK;
                            if (w->_left)
                            {
                                w->_left->_colour = BLACK;
                            }
                            RotateR(parent);
                        }
                        x = _root;
                        break;
                    }
                }
            }

            if (x)
            {
                x->_colour = BLACK;
            }
        }

        Node *Minimum(Node *node)
        {
            while (node->_left != nullptr)
            {
                node = node->_left;
            }
            return node;
        }

        void Transplant(Node *u, Node *v)
        {
            if (u->_parent == nullptr)
            {
                _root = v;
            }
            else if (u == u->_parent->_left)
            {
                u->_parent->_left = v;
            }
            else
            {
                u->_parent->_right = v;
            }

            if (v != nullptr)
            {
                v->_parent = u->_parent;
            }
        }

        // 查找节点
        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;
        }

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

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

        // void PreOrder()
        // {
        //     std::cout << "前序遍历: ";
        //     _PreOrder(_root);
        //     std::cout << std::endl;
        // }

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

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

            std::cout << "前序遍历: ";
            while (!stack.empty())
            {
                Node *cur = stack.top();
                stack.pop();
                // 访问根节点
                std::cout << cur->_kv.first << " (" << (cur->_colour == RED ? "RED" : "BLACK") << ") ";

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

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

        //     _InOrder(root->_left);
        //     // 访问根节点
        //     std::cout << root->_kv.first << " (" << (root->_colour == RED ? "RED" : "BLACK") << ") ";
        //     _InOrder(root->_right);
        // }

        // void InOrder()
        // {
        //     std::cout << "中序遍历: ";
        //     _InOrder(_root);
        //     std::cout << std::endl;
        // }

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

            std::cout << "中序遍历: ";
            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->_colour == RED ? "RED" : "BLACK") << ") ";

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

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

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

        // void PostOrder()
        // {
        //     std::cout << "后序遍历: ";
        //     _PostOrder(_root);
        //     std::cout << std::endl;
        // }

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

            std::stack<Node *> st;
            Node *cur = _root;
            Node *lastNode = nullptr; // 最近访问的节点
            std::cout << "后序遍历: ";
            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->_colour == RED ? "RED" : "BLACK") << ") ";
                    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);

            std::cout << "层序遍历: " << std::endl;
            while (!queue.empty())
            {
                int size = queue.size();
                while (size--)
                {
                    Node *cur = queue.front();
                    queue.pop();
                    // 访问当前节点
                    std::cout << cur->_kv.first << " (" << (cur->_colour == RED ? "RED" : "BLACK") << ") ";

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

        // 判断是否是红黑树
        bool IsRBTree()
        {
            // 空树也是红黑树
            if (_root == nullptr)
            {
                return true;
            }

            // 1. 根是黑色
            if (_root->_colour != BLACK)
            {
                std::cout << "根节点不是黑色" << std::endl;
                return false;
            }

            // 获取任意一条路径上黑色节点的数量
            size_t blackCount = 0;
            Node *cur = _root;
            while (cur)
            {
                if (cur->_colour == BLACK)
                {
                    blackCount++;
                }
                cur = cur->_left;
            }

            // 判断是否满足红黑树的性质, k 用来记录路径中黑色节点的个数
            size_t k = 0;
            return _IsRBTree(_root, k, blackCount);
        }

        bool _IsRBTree(Node *root, size_t k, size_t blackCount)
        {
            // 走到 nullptr 之后, 判断 k 和 blackCount 是否相等
            if (root == nullptr)
            {
                // 最终黑色节点个数
                if (blackCount != k)
                {
                    std::cout << "违反性质四: 每条路径中黑色节点的个数必须相等" << std::endl;
                    return false;
                }
                return true;
            }

            // 统计黑色节点个数
            if (root->_colour == BLACK)
            {
                k++;
            }

            // 检测当前节点与其父亲节点是否都为红色
            Node *parent = root->_parent;
            if (parent && parent->_colour == RED && root->_colour == RED)
            {
                std::cout << "违反了性质三: 红色节点的孩子必须是黑色" << std::endl;
                return false;
            }

            return _IsRBTree(root->_left, k, blackCount) && _IsRBTree(root->_right, k, blackCount);
        }

    private:
        Node *_root;
    };
}