红黑树--C++实现

#pragma once

#include <iostream>
#include <iomanip>
using namespace std;

enum RBTreeColor
{
    RED,
    BLACK
};

template <typename T>
struct RBTreeNode
{
    RBTreeNode(T value, RBTreeColor c)
        : key(value), color(c), left(nullptr), right(nullptr), parent(nullptr) {}

    RBTreeNode(T value, RBTreeColor c, RBTreeNode* l, RBTreeNode* r, RBTreeNode* p)
        : key(value), color(c), left(l), right(r), parent(p) {}

    T key;
    RBTreeColor color; 
    RBTreeNode* left;  
    RBTreeNode* right;
    RBTreeNode* parent;
};

template <typename T>
class RBTree {
public:
    typedef RBTreeNode<T> Node;
public:
    RBTree() :rootTree(nullptr) {}
    ~RBTree() { clear(); }

//查找
    //前序遍历
    void preOrder() { preOrder(rootTree); }
    //中序遍历
    void inOrder() { inOrder(rootTree); }
    //后序遍历
    void postOrder() { postOrder(rootTree); }
    //最大值
    T max();
    //最小值
    T min();
    //打印
    void print();
//修改
    //插入
    void insert(T key);
    //删除
    void remove(T key);
    // (递归实现)查找"二叉树"中键值为key的节点
    Node* search(T key);
    //清除
    void clear();
private:
    //插入
    void insert(Node* &root, Node* node);
    //插入修正
    void insertFix(Node* &root, Node* node);
    //左旋
    void leftRotate(Node* &root, Node* x);
    //右旋
    void rightRotate(Node* &root, Node* y);
    //删除
    void remove(Node* &root, Node* node);
    //删除修正
    void removeFix(Node* &root, Node* node, Node* parent);
    // (递归实现)查找"二叉树"中键值为key的节点
    Node* search(Node* tree, T key);

    void preOrder(Node *tree);
    void inOrder(Node *tree);
    void postOrder(Node *tree);
    Node* max(Node *tree);
    Node* min(Node *tree);

    void destroy(Node* &tree);

    void print(Node* tree, T key, int direction);
private:
    Node* rootTree;
};

//插入
template <typename T>
void RBTree<T>::insert(T key)
{
    Node* z = nullptr;

    if ((z = new Node(key, BLACK)) == nullptr)
        return;

    insert(rootTree, z);

}

//插入
template <typename T>
void RBTree<T>::insert(Node* &root, Node* node)
{
    Node* y = nullptr;
    Node* x = root;

    //1. 当做普通二叉查找树，将节点插入
    while (x != nullptr) {
        y = x;
        if (node->key < y->key) {
            x = y->left;
        } else {
            x = y->right;
        }
    }

    node->parent = y;
    if (y != nullptr) {
        if (node->key < y->key) {
            y->left = node;
        } else {
            y->right = node;
        }
    } else {
        root = node;
    }

    //2. 设置节点为红色
    node->color = RED;

    //3. 颜色修正
    insertFix(root, node);
}

//插入修正
template <typename T>
void RBTree<T>::insertFix(Node* &root, Node* node)
{
    Node *parent, *gparent;

    while ((parent = node->parent) && (parent->color == RED)) {
        gparent = parent->parent;

        //父节点是祖父节点的左孩子
        if (parent == gparent->left) {
            
            //case 1: 叔叔节点是红色
            {
                Node* uncle = gparent->right;
                if (uncle && (uncle->color == RED)) {
                    uncle->color = BLACK;
                    parent->color = BLACK;
                    gparent->color = RED;
                    node = gparent;
                    continue;
                }
            }

            //叔叔节点是黑色，通过旋转case 2可以转为case 3
            {
                //case 2: 叔叔节点是黑色，且当前节点是右孩子
                if (parent->right == node) {
                    Node* temp;
                    leftRotate(root, parent);
                    temp = parent;
                    parent = node;
                    node = temp;
                    //到这里情况会转为case 3
                }

                //case 3: 叔叔节点是黑色，且当前节点是左孩子。
                parent->color = BLACK;
                gparent->color = RED;
                rightRotate(root, gparent);
            }

        } else {
            
            //case 1: 叔叔节点是红色
            {
                Node* uncle = gparent->left;
                if (uncle && (uncle->color == RED)) {
                    uncle->color = BLACK;
                    parent->color = BLACK;
                    gparent->color = RED;
                    node = gparent;
                    continue;
                }
            }

            //叔叔节点是黑色，通过旋转case 2可以转为case 3
            {
                //case 2: 叔叔节点是黑色，且当前节点是右孩子
                if (parent->left == node) {
                    Node* temp;
                    rightRotate(root, parent);
                    temp = parent;
                    parent = node;
                    node = temp;
                    //到这里情况会转为case 3
                }

                //case 3: 叔叔节点是黑色，且当前节点是左孩子。
                parent->color = BLACK;
                gparent->color = RED;
                leftRotate(root, gparent);
            }
        }
    }

    //设置根节点为黑色
    root->color = BLACK;
}

/*
* 对红黑树的节点(x)进行左旋转
*
* 左旋示意图(对节点x进行左旋)：
*      px                              px
*     /                               /
*    x                               y
*   /        --(左旋)-->           /                 #
*  lx   y                          x  ry
*     /                          /  
*    ly   ry                     lx  ly
*
*
*/

//左旋
template <typename T>
void RBTree<T>::leftRotate(Node* &root, Node* x)
{
    Node* y = x->right;

    //将“y的左孩子”设为“x的右孩子”
    x->right = y->left;
    if (y->left != nullptr) {
        y->left->parent = x;
    }

    y->parent = x->parent;

    if (x->parent == nullptr) {
        root = y;
    } else {
        if (x->parent->left == x) {
            x->parent->left = y;
        } else {
            x->parent->right = y;
        }
    }

    y->left = x;

    x->parent = y;
}

/*
* 对红黑树的节点(y)进行右旋转
*
* 右旋示意图(对节点y进行左旋)：
*            py                               py
*           /                                /
*          y                                x
*         /        --(右旋)-->            /                       #
*        x   ry                           lx   y
*       /                                    /                    #
*      lx  rx                                rx  ry
*
*/
//右旋
template <typename T>
void RBTree<T>::rightRotate(Node* &root, Node* y)
{
    Node *x = y->left;

    //调整y左孩子
    y->left = x->right;
    if (x->right != nullptr) {
        x->right->parent = nullptr;
    }

    //调整x节点
    x->parent = y->parent;

    if (y->parent == nullptr) {
        root = x;
    } else {
        if (y->parent->left == y) {
            y->parent->left = x;
        } else {
            y->parent->right = x;
        }
    }

    x->right = y;

    //调整y
    y->parent = x;
}

//删除
template <typename T>
void RBTree<T>::remove(T key)
{
    Node* z;
    if ((z = search(key)) != nullptr) {
        remove(rootTree, z);
    }
}

//删除
template <typename T>
void RBTree<T>::remove(Node* &root, Node* node)
{
    Node *child, *parent;
    RBTreeColor color;

    //1.删除节点：当作普通二叉树
    //2.删除后，修正颜色

    //case 1：节点有两个孩子
    if ((node->left != nullptr) && (node->right != nullptr)) {
        
        //查找取代节点
        Node* replace = node;

        //获取后继节点
        replace = replace->right;
        while (replace->left != nullptr) {
            replace = replace->left;
        }

        //移动后继节点
        if (node->parent) {
            if (node->parent->left == node) {
                node->parent->left = replace;
            } else {
                node->parent->right = replace;
            }
        } else {
            root = replace;
        }

        //后继节点有孩子一定是右孩子
        child = replace->right;
        parent = replace->parent;
        color = replace->color;

        //分两种情况，replace是否是node的子孩子
        if (parent == node) {
            //子孩子比较简单，直接删除node，replace替换上
            parent = replace;
        } else {
            //
            if (child) {
                child->parent = parent;
            }

            parent->left = child;
            replace->right = node->right;
            node->right = replace;
        }
        replace->parent = node->parent;
        replace->color = node->color;
        replace->left = node->left;
        node->left->parent = replace;

        //节点为黑色需要修正
        if (color == BLACK) {
            removeFix(root, child, parent);
        }

        delete node;
        return;
    }
    
    //case 2:含有一个孩子或者无
    if (node->left != nullptr) {
        child = node->left;
    } else {
        child = node->right;
    }

    parent = node->parent;
    color = node->color;

    if (child) {
        child->parent = parent;
    }

    if (parent) {
        if (parent->left == node) {
            parent->left = child;
        } else {
            parent->right = child;
        } 
    } else {
        root = child;
    }

    if (color == BLACK) {
        removeFix(root, child, parent);
    }
    delete node;
}

//删除修正
template <typename T>
void RBTree<T>::removeFix(Node* &root, Node* node, Node* parent)
{
    Node* other;

    while ((!node || node->color == BLACK) && node != root)
    {
        if (parent->left == node) {
            other = parent->right;

            //case 1:兄弟是红色 ==> 调整为case 2 3 4
            //此时parent任然是不平衡的
            if (other->color == RED) {
                other->color = BLACK;
                parent->color = RED;
                leftRotate(root, parent);
                other = parent->right;
            }

            //case 2:兄弟w是黑色，且w两个孩子也是黑色
            if ((!other->left || other->left->color == BLACK) &&
                (!other->right || other->right->color == BLACK)) {
                other->color = RED;
                node = parent;
                parent = node->parent;
            } else {

                //case 3:兄弟w是黑色，且w的左孩子是红色，右孩子是黑色 ==> 调整为case 4
                if (!other->right || other->right->color == BLACK) {
                    other->left->color = BLACK;
                    other->color = RED;
                    rightRotate(root, other);
                    other = parent->right;
                }

                //case 4:兄弟w是黑色，且w的右孩子为红色
                other->color = parent->color;
                parent->color = BLACK;
                other->right->color = BLACK;
                leftRotate(root, parent);
                node = root;
                break;
            }
        } else {
            //右节点，情况相反
            other = parent->left;

            //case 1:兄弟是红色 ==> 调整为case 2 3 4
            //此时parent任然是不平衡的
            if (other->color == RED) {
                other->color = BLACK;
                parent->color = RED;
                leftRotate(root, parent);
                other = parent->left;
            }

            //case 2:兄弟w是黑色，且w两个孩子也是黑色
            if ((!other->left || other->left->color == BLACK) &&
                (!other->right || other->right->color == BLACK)) {
                other->color = RED;
                node = parent;
                parent = node->parent;
            } else {
                //case 3:兄弟w是黑色，且w的左孩子是红色，右孩子是黑色 ==> 调整为case 4
                if (!other->left || other->left->color == BLACK) {
                    other->right->color = BLACK;
                    other->color = RED;
                    rightRotate(root, other);
                    other = parent->left;
                }

                //case 4:兄弟w是黑色，且w的右孩子为红色
                other->color = parent->color;
                parent->color = BLACK;
                other->left->color = BLACK;
                leftRotate(root, parent);
                node = root;
                break;
            }
        }
    }

    if (node) {
        node->color = BLACK;
    }
}

template <typename T>
typename RBTree<T>::Node* RBTree<T>::search(Node* tree, T key)
{

    if (tree == nullptr || tree->key == key) {
        return tree;
    }

    if (key < tree->key) {
        return search(tree->left, key);
    }
    else {
        return search(tree->right, key);
    }
}

template <typename T>
typename RBTree<T>::Node* RBTree<T>::search(T key)
{
    return search(rootTree, key);
}


template <typename T>
typename RBTree<T>::Node*  RBTree<T>::max(Node *tree)
{
    if (!tree) return nullptr;

    Node* x = tree;
    while (x->right) {
        x = x->right;
    }
    return x;
}

template <typename T>
T RBTree<T>::max()
{
    Node* x = max(rootTree);

    if (x) {
        return x->key;
    }
    else {
        return (T)nullptr;
    }

}

template <typename T>
typename RBTree<T>::Node* RBTree<T>::min(Node *tree)
{
    if (!tree) return nullptr;

    Node* x = tree;
    while (x->left) {
        x = x->left;
    }
    return x;
}

template <typename T>
T RBTree<T>::min()
{
    Node* x = min(rootTree);

    if (x) {
        return x->key;
    }
    else {
        return (T)nullptr;
    }
}

template <typename T>
void RBTree<T>::preOrder(Node *tree)
{
    if (tree == nullptr) return;
    std::cout << tree->key << " ";
    preOrder(tree->left);
    preOrder(tree->right);
}

template <typename T>
void RBTree<T>::inOrder(Node *tree)
{
    if (!tree) return;
    inOrder(tree->left);
    std::cout << tree->key << " ";
    inOrder(tree->right);
}

template <typename T>
void RBTree<T>::postOrder(Node *tree)
{
    if (!tree) return;
    postOrder(tree->left);
    postOrder(tree->right);
    std::cout << tree->key << " ";
}

template <typename T>
void RBTree<T>::clear()
{
    std::cout << "===destroy: ";
    destroy(rootTree);
    std::cout << std::endl;
}

template <typename T>
void RBTree<T>::destroy(Node* &tree)
{
    if (tree == nullptr) return;

    if (tree->left) {
        destroy(tree->left);
    }
    if (tree->right) {
        destroy(tree->right);
    }

    std::cout << tree->key << " ";
    delete tree;
    tree = nullptr;
}

/*
* 打印"二叉查找树"
*
* key        -- 节点的键值
* direction  --  0，表示该节点是根节点;
*               -1，表示该节点是它的父结点的左孩子;
*                1，表示该节点是它的父结点的右孩子。
*/
template < typename T>
void RBTree<T>::print(Node* tree, T key, int direction)
{
     if (tree != nullptr)
    {
         if (direction == 0)    // tree是根节点
             cout << setw(2) << tree->key << "(B) is root" << endl;
         else                // tree是分支节点
             cout << setw(2) << tree->key << ( (tree->color == RED) ? "(R)" : "(B)" )<< " is " << setw(2) << key << "'s " << setw(12) << (direction == 1 ? "right child" : "left child") << endl;

         print(tree->left, tree->key, -1);
         print(tree->right, tree->key, 1);
     }
}


template < typename T>
void RBTree<T>::print()
{
    if (rootTree != NULL)
         print(rootTree, rootTree->key, 0);
}