Set_Map/Linux/RBTree.hpp

#pragma once

#include <iostream>

namespace Lenyiin
{
    enum Colour
    {
        RED,
        BLACK
    };

    template <class T>
    struct RBTreeNode
    {
        RBTreeNode<T> *_left;   // 左子节点
        RBTreeNode<T> *_right;  // 右子节点
        RBTreeNode<T> *_parent; // 父节点

        T _data;        // 节点存储的数据
        Colour _colour; // 节点的颜色（红色或黑色）

        // 节点的构造函数，默认为红色节点
        RBTreeNode(const T &data)
            : _left(nullptr), _right(nullptr), _parent(nullptr), _data(data), _colour(RED)
        {
        }
    };

    template <class T, class Ref, class Ptr>
    struct __TreeIterator
    {
        typedef RBTreeNode<T> Node;
        typedef __TreeIterator<T, Ref, Ptr> Self;
        Node *_node; // 当前迭代器指向的节点

        // 构造函数
        __TreeIterator(Node *node)
            : _node(node)
        {
        }

        // 解引用操作符
        Ref operator*()
        {
            return _node->_data;
        }

        // 访问成员操作符
        Ptr operator->()
        {
            return &_node->_data;
        }

        // 前置递增运算符
        Self &operator++()
        {
            // 1. 如果右不为空, 中序的下一个就是右子树的最左节点
            if (_node->_right)
            {
                Node *subLeft = _node->_right;
                while (subLeft->_left)
                {
                    subLeft = subLeft->_left;
                }
                _node = subLeft;
            }
            // 2. 如果右为空, 表示 _node 所在的子树已经完成, 下一个节点在他祖先中去找, 沿着路径往上找孩子使它的左的那个祖先
            else
            {
                Node *cur = _node;
                Node *parent = cur->_parent;
                while (parent && cur == parent->_right)
                {
                    cur = parent;
                    parent = cur->_parent;
                }
                _node = parent;
            }
            return *this;
        }

        // 后置递增运算符
        Self operator++(int)
        {
            Node *tmp = _node;
            ++(*this);
            return Self(tmp);
        }

        // 前置递减运算符
        Self &operator--()
        {
            // 1. 如果左不为空, 中序的上一个就是左树的最右节点
            if (_node->_left)
            {
                Node *subRight = _node->left;
                while (subRight->_right)
                {
                    subRight = subRight->_right;
                }
                _node = subRight;
            }
            // 2. 如果左为空, 表示 _node 所在的子树已经完成, 上一个节点在他祖先中去找, 沿着路径往上找孩子是它的右的那个祖先
            else
            {
                Node *cur = _node;
                Node *parent = cur->_parent;
                while (parent && cur == parent->_left)
                {
                    cur = parent;
                    parent = cur->_parent;
                }
                _node = parent;
            }
            return *this;
        }

        // 后置递减运算符
        Self operator--(int)
        {
            Node *tmp = _node;
            --(*this);
            return Self(tmp);
        }

        bool operator!=(const Self &n)
        {
            return _node != n._node;
        }

        bool operator==(const Self &n)
        {
            return _node == n._node;
        }
    };

    template <class K, class T, class KOfT>
    class RBTree
    {
    private:
        typedef RBTreeNode<T> 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 (_root == parent)
            {
                _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 (_root == parent)
            {
                _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:
        typedef __TreeIterator<T, T &, T *> iterator;
        typedef __TreeIterator<T, const T &, const T *> const_iterator;

        iterator begin()
        {
            Node *cur = _root;
            while (cur && cur->_left)
            {
                cur = cur->_left;
            }
            return iterator(cur);
        }

        iterator end()
        {
            return iterator(nullptr);
        }

        const_iterator begin() const
        {
            Node *cur = _root;
            while (cur && cur->_left)
            {
                cur = cur->_left;
            }
            return const_iterator(cur);
        }

        const_iterator end() const
        {
            return const_iterator(nullptr);
        }

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

        ~RBTree()
        {
            DeleteTree(_root);
        }

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

            // 插入时使用比较器来比较键的大小，以支持不同的数据类型
            KOfT koft;
            Node *parent = nullptr;
            Node *cur = _root;
            while (cur)
            {
                if (koft(cur->_data) > koft(data))
                {
                    parent = cur;
                    cur = cur->_left;
                }
                else if (koft(cur->_data) < koft(data))
                {
                    parent = cur;
                    cur = cur->_right;
                }
                else
                {
                    // 如果键已经存在, 插入无效, 返回 false, 并且返回该键的迭代器
                    return std::make_pair(iterator(cur), false);
                }
            }

            // 找到位置, 根据父节点插入新节点
            cur = new Node(data);
            Node *newnode = cur;
            if (koft(parent->_data) > koft(cur->_data))
            {
                parent->_left = cur;
                cur->_parent = parent;
            }
            else
            {
                parent->_right = cur;
                cur->_parent = parent;
            }

            // 插入后需要修复红黑树的性质
            InsertFixUp(parent, cur);

            return std::make_pair(iterator(newnode), 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 = 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;
            KOfT koft;

            // 1. 寻找要删除的节点
            while (nodeToDelete)
            {
                if (koft(nodeToDelete->_data) > key)
                {
                    nodeToDelete = nodeToDelete->_left;
                }
                else if (koft(nodeToDelete->_data) < key)
                {
                    nodeToDelete = nodeToDelete->_right;
                }
                else
                {
                    break; // 找到了要删除的节点
                }
            }

            // 节点若不存在, 直接返回 false
            if (nodeToDelete == nullptr)
            {
                return false;
            }

            // 执行删除操作
            Node *parent, *child;
            // 保存原节点的颜色，以便后续调整
            Colour originalColour = nodeToDelete->_colour;

            // 2. 处理删除节点的各种情况
            if (nodeToDelete->_left == nullptr)
            {
                // 情况 1：没有左子节点
                child = nodeToDelete->_right;
                parent = nodeToDelete->_parent;
                Transplant(nodeToDelete, nodeToDelete->_right);
            }
            else if (nodeToDelete->_right == nullptr)
            {
                // 情况 2：没有右子节点
                child = nodeToDelete->_left;
                parent = nodeToDelete->_parent;
                Transplant(nodeToDelete, nodeToDelete->_left);
            }
            else
            {
                // 情况 3：有两个子节点
                // 找到右子树中的最小节点（后继节点）
                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;

            // 3. 修复红黑树的性质
            // 如果删除的节点是黑色, 需要进行调整
            if (originalColour == BLACK)
            {
                EraseFixUp(child, parent);
            }

            return true;
        }

        Node *Minimum(Node *node)
        {
            while (node->_left != nullptr)
            {
                node = node->_left; // 一直向左走，直到找到最左节点
            }
            return node;
        }

        void Transplant(Node *u, Node *v)
        {
            // 如果u是根节点，v成为新的根节点
            if (u->_parent == nullptr)
            {
                _root = v;
            }
            // u是左子节点，用v替换它
            else if (u == u->_parent->_left)
            {
                u->_parent->_left = v;
            }
            // u是右子节点，用v替换它
            else
            {
                u->_parent->_right = v;
            }

            // 连接v与u的父节点
            if (v != nullptr)
            {
                v->_parent = u->_parent;
            }
        }

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

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

        // 查找节点
        iterator &Find(const K &key)
        {
            KOfT koft;
            Node *cur = _root;
            while (cur)
            {
                if (koft(cur->_data) > key)
                {
                    cur = cur->_left;
                }
                else if (koft(cur->_data) < key)
                {
                    cur = cur->_right;
                }
                else
                {
                    return iterator(cur);
                }
            }
            return iterator(nullptr);
        }

        // 判断是否是红黑树
        bool IsRBTree()
        {
            Node *root = _root;

            // 空树也是红黑树
            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;
    };
}