红黑树的实现（C++）

本文具体实现了红黑树的基本操作，红黑树的相关定义和性质请移步到 红黑树（算法导论版）_青衫客36的博客-CSDN博客

#include <iostream>
using namespace std;

//  1  结点定义：rbTreeNode
template <class T>
class rbtree;

template <class T>
struct _rbTreeNode
{
    friend class rbtree<T>;
    T getkey()
    {
        return key;
    }

private:
    T key;
    _rbTreeNode<T> *left;
    _rbTreeNode<T> *right;
    _rbTreeNode<T> *p;
    bool color; // true == red  &&  false == black
};

//  2  红黑树类声明：rbtree
template <class T>
class rbtree
{
public:
    rbtree() : nil(new _rbTreeNode<T>), root(nil)
    {
        nil->left = nullptr;
        nil->p = nullptr;
        nil->right = nullptr;
        nil->color = false;
    }
    void insert(T key);
    _rbTreeNode<T> *search(T key);
    void erase(T key);
    bool isRBTree();
    void display();

private:
    void _display(_rbTreeNode<T> *);

    bool _isRBTree(_rbTreeNode<T> *, int count, int blackCount);

    void leftRotate(_rbTreeNode<T> *);
    void rightRotate(_rbTreeNode<T> *);

    void rbDelete(_rbTreeNode<T> *);
    void rbDeleteFixup(_rbTreeNode<T> *);

    void rbInsertFixup(_rbTreeNode<T> *);

    _rbTreeNode<T> *treeSuccessor(_rbTreeNode<T> *);

    _rbTreeNode<T> *nil; // 空结点
    _rbTreeNode<T> *root;
};

//  3  红黑树类成员定义
template <class T>
void rbtree<T>::insert(T key)
{
    _rbTreeNode<T> *t = new _rbTreeNode<T>;
    t->key = key;
    _rbTreeNode<T> *x = root;
    _rbTreeNode<T> *y = nil;
    while (x != nil) // 找到应插入的位置
    {
        y = x;
        if (key < x->key)
            x = x->left;
        else
            x = x->right;
    }
    t->p = y;
    if (y == nil)
        root = t;
    else // 将新插入的节点t经判断后插在这个位置
    {
        if (t->key < y->key)
            y->left = t;
        else
            y->right = t;
    }
    t->left = nil;
    t->right = nil;
    t->color = true;
    rbInsertFixup(t);
}

template <class T>
_rbTreeNode<T> *rbtree<T>::search(T key)
{
    _rbTreeNode<T> *x = root;
    while (x != nil && key != x->key)
    {
        if (key < x->key)
            x = x->left;
        else
            x = x->right;
    }
    return x;
}

template <class T>
void rbtree<T>::erase(T key)
{
    _rbTreeNode<T> *x = search(key);
    if (x != nil)
        rbDelete(x);
}

template <class T>
_rbTreeNode<T> *rbtree<T>::treeSuccessor(_rbTreeNode<T> *x)
{
    x = x->right;
    while (x->left != nil)
        x = x->left;
    return x;
    // if (x->right != nil)
    // {
    //     x = x->right;
    //     while (x->left != nil)
    //         x = x->left;
    //     return x;
    // }
    // _rbTreeNode<T> *y = x->p;
    // while (y != nil && x == y->right)
    // {
    //     x = y;
    //     y = y->p;
    // }
    // return y;
}

template <class T>
void rbtree<T>::leftRotate(_rbTreeNode<T> *x)
{ // x->right != nil
    _rbTreeNode<T> *y = x->right;
    if (x->p == nil)
        root = y;
    else
    {
        if (x == x->p->left)
            x->p->left = y;
        else
            x->p->right = y;
    }
    y->p = x->p;
    x->right = y->left;
    y->left->p = x;
    y->left = x;
    x->p = y;
}

template <class T>
void rbtree<T>::rightRotate(_rbTreeNode<T> *x)
{
    _rbTreeNode<T> *y = x->left;
    if (x->p == nil)
        root = y;
    else
    {
        if (x->p->left == x)
            x->p->left = y;
        else
            x->p->right = y;
    }
    y->p = x->p;
    x->left = y->right;
    y->right->p = x;
    y->right = x;
    x->p = y;
}

template <class T>
void rbtree<T>::rbInsertFixup(_rbTreeNode<T> *z)
{
    while (z->p->color == true) // 被插入的节点的父节点是红色
    {
        if (z->p == z->p->p->left) // 父节点是爷节点的左孩子
        {
            _rbTreeNode<T> *y = z->p->p->right;
            if (y->color == true) // 当前节点的叔节点也是红色
            {
                z->p->color = false;   // 将父节点设为黑色
                y->color = false;      // 将叔节点设为黑色
                z->p->p->color = true; // 将祖父节点设为红色
                z = z->p->p;           // 将祖父节点设为当前节点（红色节点）；之后继续对当前节点进行操作
            }
            else // 当前节点的叔节点是黑色
            {
                if (z == z->p->right) // 当前节点是其父节点的右孩子
                {
                    z = z->p;      // 将父节点作为新的当前节点
                    leftRotate(z); // 以新的当前节点为支点进行左旋
                }

                // 当前节点是其父节点的左孩子
                z->p->color = false;   // 将父节点设为黑色
                z->p->p->color = true; // 将祖父节点设为红色
                // z->p->right->color = false;
                rightRotate(z->p->p); // 以祖父节点为支点进行右旋
            }
        }
        else // // 父节点是爷节点的右孩子
        {
            _rbTreeNode<T> *y = z->p->p->left;
            if (y->color == true)
            {
                z->p->color = false;
                y->color = false;
                z->p->p->color = true;
                z = z->p->p;
            }
            else
            {
                if (z == z->p->left)
                {
                    z = z->p;
                    rightRotate(z);
                }
                z->p->color = false;
                z->p->p->color = true;
                // z->p->left->color = false;
                leftRotate(z->p->p);
            }
        }
    }
    root->color = false; // 根节点设为黑色
}

template <class T>
void rbtree<T>::rbDelete(_rbTreeNode<T> *z) // z是逻辑上被删除的节点(实际上他并没有被删除，而是修改了自己的value值为中序后继的值，然后把中序后继节点y删除了)
{
    _rbTreeNode<T> *x = nil;               // 取代被删除节点位置的节点
    _rbTreeNode<T> *y = nil;               // 真正删除的节点
    if (z->left == nil || z->right == nil) // z只有左子树或只有右子树，就直接把子树的根节点提上去
        y = z;
    else
        y = treeSuccessor(z); // z左右子树都有，那么就把z的中序后继提上去
    if (y->left != nil)       // 然后准备删除y，类似于链表的删除，将y的子节点挂在y的父节点下(取代y的位置)
        x = y->left;
    else
        x = y->right;
    x->p = y->p;
    if (y->p == nil)
        root = x;
    else
    {
        if (y == y->p->left)
            y->p->left = x;
        else
            y->p->right = x; // 将y的父节点与x进行双向链接
    }
    if (y != z)            // 逻辑上删除的节点与真正删除的节点不是同一个，就把实际上删除的节点的值赋给逻辑上删除的节点，(是同一个节点的话，直接把该节点删了就可以了)
        z->key = y->key;   // 将y的值赋给z使之满足二叉搜索树的性质
    if (y->color == false) // 如果被删除的节点是黑色，需要对红黑树的特性进行调整
        rbDeleteFixup(x);
    delete y; // 如果删除的节点是红色，直接删除就可以了，不需要调整红黑树
}

template <class T>
void rbtree<T>::rbDeleteFixup(_rbTreeNode<T> *x)
{
    // 如果x是红色节点，不用执行该while循环，直接将其设为黑色即可
    // 如果x是黑色节点，需要对红黑树进行调整，以满足红黑树的特性
    while (x != root && x->color == false)
    {
        _rbTreeNode<T> *w = 0;
        if (x->p->left == x) // x是父节点的左孩子
        {
            w = x->p->right;
            if (w->color == true) // x的兄弟节点是红色
            {
                w->color = false;   // 将x的兄弟节点设为黑色
                x->p->color = true; // 将x的父节点设为红色
                leftRotate(x->p);   // 对x的父节点进行左旋
                w = x->p->right;    // 左旋后，重新设置x的兄弟节点
            }
            // x的兄弟节点是黑色，x的兄弟节点的两个孩子都是黑色
            if (w->left->color == false && w->right->color == false)
            {
                w->color = true; // 将x的兄弟节点设为红色
                x = x->p;        // 设置x的父节点为新的x节点
            }
            else
            {
                // x的兄弟节点是黑色，x的兄弟节点的左孩子是红色，右孩子是黑色
                if (w->right->color == false)
                {
                    w->left->color = false; // 将x兄弟节点的左孩子设为黑色
                    w->color = true;        // 将x兄弟节点设为红色
                    rightRotate(w);         // 对x的兄弟节点进行右旋
                    w = x->p->right;        // 右旋后，重新设置x的兄弟节点
                }

                // x的兄弟节点是黑色，x的兄弟节点的右孩子是红色的，x的兄弟节点的左孩子任意颜色
                w->color = x->p->color;  // 将x父节点颜色赋值给x的兄弟节点
                x->p->color = false;     // 将x父节点设置为黑色
                w->right->color = false; // 将x兄弟节点的右子节点设为黑色
                leftRotate(x->p);        // 对x的父节点进行左旋
                x = root;
            }
        }
        else // x是父节点的右孩子
        {
            w = x->p->left;
            if (w->color == true)
            {
                w->color = false;
                x->p->color = true;
                rightRotate(x->p);
                w = x->p->left;
            }
            if (w->right->color == false && w->left->color == false)
            {
                w->color = true;
                x = x->p;
            }
            else
            {
                if (w->left->color == false)
                {
                    w->right->color = false;
                    w->color = true;
                    leftRotate(w);
                    w = x->p->left;
                }
                w->color = x->p->color;
                x->p->color = false;
                w->left->color = false;
                rightRotate(x->p);
                x = root;
            }
        }
    }
    x->color = false; // 将根节点设为黑色
}

template <class T>
void rbtree<T>::display()
{
    if (root != nil)
        _display(root);
    else
        cout << "Tree is empty! " << endl;
}

template <class T>
bool rbtree<T>::isRBTree()
{
    if (root == nullptr)
        return true;
    if (root->color == true)
    {
        cout << "error: 根节点是红色" << endl;
        return false;
    }

    // 找最左路径作为黑色节点数目的参考值
    _rbTreeNode<T> *cur = root;
    int blackCount = 0;
    while (cur)
    {
        if (cur->color == false)
            blackCount++;
        cur = cur->left;
    }

    int count = 0;
    return _isRBTree(root, count, blackCount);
}

template <class T>
bool rbtree<T>::_isRBTree(_rbTreeNode<T> *root, int count, int blackCount)
{
    if (root == nullptr)
    {
        if (count != blackCount)
        {
            cout << "error: 黑色节点的数目不相等" << endl;
            return false;
        }
        return true;
    }

    if (root->color == true && root->p->color == true)
    {
        cout << "error: 存在连续的红色节点" << endl;
        return false;
    }

    if (root->color == false)
        count++;

    return _isRBTree(root->left, count, blackCount) && _isRBTree(root->right, count, blackCount);
}

template <class T>
void rbtree<T>::_display(_rbTreeNode<T> *x)
{
    if (x->left != nil)
        _display(x->left);
    if (x != nil)
    {

        if (x->color == true)
            cout << "RED   ";
        else
            cout << "BLACK ";

        cout << x->key << ' ';
    }
    cout << endl;
    if (x->right != nil)
        _display(x->right);
}

void judge(bool flag)
{
    if (flag)
        cout << "It is a Red-Black Tree" << endl;
    else
        cout << "It is not a Red-Black Tree" << endl;
}

//  4  测试程序：main()
// const int N = 12;

int p[12] = {1, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15};
// int p[6] = {41, 38, 31, 12, 19, 8};

int main()
{
    rbtree<int> test;
    for (int i = 0; i < 12; ++i)
    {
        test.insert(p[i]);
    }
    test.display();
    judge(test.isRBTree());

    int a;
    while (cin >> a)
    {
        test.erase(a);
        // cout << test.isRBTree() << endl;
        test.display();
        judge(test.isRBTree());
    }
    return 0;
}