/*** C++ 语言: 红黑树** @author skywang* @date 2013/11/07** @update fyh* @date 2020/09/22*/
#pragma once#include <iomanip>
#include <iostream>
using namespace std;/* 节点颜色 */
enum rbt_color
{rbt_red,			// 红色rbt_black			// 黑色
};/* 节点类 */
template <class T>
class rbt_node
{
public:rbt_color m_color = rbt_red;				// 颜色T m_key;												// 关键字(键值)rbt_node *m_left = nullptr;					// 左孩子rbt_node *m_right = nullptr;				// 右孩子rbt_node *m_parent = nullptr;			// 父结点rbt_node(T value, rbt_color c, rbt_node *p, rbt_node *l, rbt_node *r) :m_key(value), m_color(c), m_parent(), m_left(l), m_right(r) {}
};/* 树类 */
template <class T>
class rbt_tree
{
private:rbt_node<T> *m_root = nullptr;    // 根结点public:/* 构造函数 */rbt_tree();/* 析构函数 */~rbt_tree();// 前序遍历"红黑树"void pre_order();// 中序遍历"红黑树"void in_order();// 后序遍历"红黑树"void post_order();// (递归实现)查找"红黑树"中键值为key的节点rbt_node<T>* search(T key);// (非递归实现)查找"红黑树"中键值为key的节点rbt_node<T>* iterative_search(T key);// 查找最小结点：返回最小结点的键值。T minimum();// 查找最大结点：返回最大结点的键值。T maximum();// 找结点(x)的后继结点。即，查找"红黑树中数据值大于该结点"的"最小结点"。rbt_node<T>* successor(rbt_node<T> *x);// 找结点(x)的前驱结点。即，查找"红黑树中数据值小于该结点"的"最大结点"。rbt_node<T>* predecessor(rbt_node<T> *x);// 将结点(key为节点键值)插入到红黑树中void insert(T key);// 删除结点(key为节点键值)void remove(T key);// 销毁红黑树void destroy();// 打印红黑树void print();private:// 错误处理void error(bool state, const char* text) const;// 前序遍历"红黑树"void pre_order(rbt_node<T>* tree) const;// 中序遍历"红黑树"void in_order(rbt_node<T>* tree) const;// 后序遍历"红黑树"void post_order(rbt_node<T>* tree) const;// (递归实现)查找"红黑树x"中键值为key的节点rbt_node<T>* search(rbt_node<T>* x, T key) const;// (非递归实现)查找"红黑树x"中键值为key的节点rbt_node<T>* iterative_search(rbt_node<T>* x, T key) const;// 查找最小结点：返回tree为根结点的红黑树的最小结点。rbt_node<T>* minimum(rbt_node<T>* tree);// 查找最大结点：返回tree为根结点的红黑树的最大结点。rbt_node<T>* maximum(rbt_node<T>* tree);// 左旋void left_rotate(rbt_node<T>* &root, rbt_node<T>* x);// 右旋void right_rotate(rbt_node<T>* &root, rbt_node<T>* y);// 插入函数void insert(rbt_node<T>* &root, rbt_node<T>* node);// 插入修正函数void insert_fix_up(rbt_node<T>* &root, rbt_node<T>* node);// 删除函数void remove(rbt_node<T>* &root, rbt_node<T> *node);// 删除修正函数void remove_fix_up(rbt_node<T>* &root, rbt_node<T> *node, rbt_node<T> *parent);// 销毁红黑树void destroy(rbt_node<T>* &tree);// 打印红黑树void print(rbt_node<T>* tree, T key, int direction);inline rbt_node<T>* rb_parent(rbt_node<T>* r) { return r->m_parent; }inline rbt_color rb_color(rbt_node<T>* r) { return r->m_color; }inline bool rb_is_red(rbt_node<T>* r) { return r->m_color == rbt_red; }inline bool rb_is_black(rbt_node<T>* r) { return r->m_color == rbt_black; }inline void rb_set_black(rbt_node<T>* r) { r->m_color = rbt_black; }inline void rb_set_red(rbt_node<T>* r) { r->m_color = rbt_red; }inline void rb_set_parent(rbt_node<T>* r, rbt_node<T>* p) { r->m_parent = p; }inline void rb_set_color(rbt_node<T>* r, rbt_color c) { r->m_color = c; }//#define rb_parent(r)   ((r)->m_parent)//#define rb_color(r) ((r)->m_color)//#define rb_is_red(r)   ((r)->color==rbt_red)//#define rb_is_black(r)  ((r)->color==rbt_black)//#define rb_set_black(r)  do { (r)->color = rbt_black; } while (0)//#define rb_set_red(r)  do { (r)->color = rbt_red; } while (0)//#define rb_set_parent(r,p)  do { (r)->m_parent = (p); } while (0)//#define rb_set_color(r,c)  do { (r)->m_color = (c); } while (0)
};/** 构造函数*/
template <class T>
rbt_tree<T>::rbt_tree() :m_root(nullptr) {}/** 析构函数*/
template <class T>
rbt_tree<T>::~rbt_tree() { destroy(); }/*
* 错误处理
*/
template <class T>
void rbt_tree<T>::error(bool state, const char* text) const
{if (state == false){std::cout << text << std::endl;getchar();exit(-1);}
}/** 前序遍历"红黑树"* 中 -> 左 -> 右*/
template <class T>
void rbt_tree<T>::pre_order(rbt_node<T>* tree) const
{if (tree){cout << tree->m_key << " ";pre_order(tree->m_left);pre_order(tree->m_right);}
}template <class T>
void rbt_tree<T>::pre_order()
{pre_order(m_root);
}/** 中序遍历"红黑树"* 左 -> 中 -> 右*/
template <class T>
void rbt_tree<T>::in_order(rbt_node<T>* tree) const
{if (tree){in_order(tree->m_left);cout << tree->m_key << " ";in_order(tree->m_right);}
}template <class T>
void rbt_tree<T>::in_order()
{in_order(m_root);
}/** 后序遍历"红黑树"* 左 -> 右 -> 中*/
template <class T>
void rbt_tree<T>::post_order(rbt_node<T>* tree) const
{if (tree){post_order(tree->m_left);post_order(tree->m_right);cout << tree->m_key << " ";}
}template <class T>
void rbt_tree<T>::post_order()
{post_order(m_root);
}/** (递归实现)查找"红黑树x"中键值为key的节点*/
template <class T>
rbt_node<T>* rbt_tree<T>::search(rbt_node<T>* x, T key) const
{// 直接返回if (x == nullptr || x->m_key == key)return x;if (key < x->m_key)return search(x->m_left, key);		// 小于搜索左边elsereturn search(x->m_right, key);	// 大于搜索右边
}template <class T>
rbt_node<T>* rbt_tree<T>::search(T key)
{return search(m_root, key);
}/** (非递归实现)查找"红黑树x"中键值为key的节点*/
template <class T>
rbt_node<T>* rbt_tree<T>::iterative_search(rbt_node<T>* x, T key) const
{// x不为空且键值不相同while (x && x->m_key != key){if (key < x->m_key)x = x->m_left;		// 小于就往左搜索elsex = x->m_right;	// 大于就往右搜索}return x;
}template <class T>
rbt_node<T>* rbt_tree<T>::iterative_search(T key)
{return iterative_search(m_root, key);
}/** 查找最小结点：返回tree为根结点的红黑树的最小结点。*/
template <class T>
rbt_node<T>* rbt_tree<T>::minimum(rbt_node<T>* tree)
{// 为空就直接返回空if (tree == nullptr)return nullptr;// 最左的那一个就是最小的那一个while (tree->m_left)tree = tree->m_left;return tree;
}template <class T>
T rbt_tree<T>::minimum()
{return minimum(m_root);
}/** 查找最大结点：返回tree为根结点的红黑树的最大结点。*/
template <class T>
rbt_node<T>* rbt_tree<T>::maximum(rbt_node<T>* tree)
{// 传入为空那就直接返回空if (tree == nullptr)return nullptr;// 最右边的那一个就是最大的那一个while (tree->m_right)tree = tree->m_right;return tree;
}template <class T>
T rbt_tree<T>::maximum()
{return maximum(m_root);
}/** 找结点(x)的后继结点。即，查找"红黑树中数据值大于该结点"的"最小结点"。*/
template <class T>
rbt_node<T>* rbt_tree<T>::successor(rbt_node<T> *x)
{// 如果x存在右孩子，则"x的后继结点"为 "以其右孩子为根的子树的最小结点"。if (x->m_right)return minimum(x->m_right);// 如果x没有右孩子。则x有以下两种可能：// (01) x是"一个左孩子"，则"x的后继结点"为 "它的父结点"。// (02) x是"一个右孩子"，则查找"x的最低的父结点，并且该父结点要具有左孩子"，找到的这个"最低的父结点"就是"x的后继结点"。rbt_node<T>* y = x->m_parent;while (y && x == y->m_right){x = y;y = y->m_parent;}return y;
}/** 找结点(x)的前驱结点。即，查找"红黑树中数据值小于该结点"的"最大结点"。*/
template <class T>
rbt_node<T>* rbt_tree<T>::predecessor(rbt_node<T> *x)
{// 如果x存在左孩子，则"x的前驱结点"为 "以其左孩子为根的子树的最大结点"。if (x->m_left)return maximum(x->m_left);// 如果x没有左孩子。则x有以下两种可能：// (01) x是"一个右孩子"，则"x的前驱结点"为 "它的父结点"。// (01) x是"一个左孩子"，则查找"x的最低的父结点，并且该父结点要具有右孩子"，找到的这个"最低的父结点"就是"x的前驱结点"。rbt_node<T>* y = x->m_parent;while (y&& x == y->m_left){x = y;y = y->m_parent;}return y;
}/** 对红黑树的节点(x)进行左旋转** 左旋示意图(对节点x进行左旋)：*      px                              px*     /                               /*    x                               y*   /  \      --(左旋)-->     / \                #*  lx   y                          x  ry*     /   \                       /  \*    ly   ry                     lx  ly***/
template <class T>
void rbt_tree<T>::left_rotate(rbt_node<T>* &root, rbt_node<T>* x)
{// 设置x的右孩子为yrbt_node<T> *y = x->m_right;// 将 “y的左孩子” 设为 “x的右孩子”；// 如果y的左孩子非空，将 “x” 设为 “y的左孩子的父亲”x->m_right = y->m_left;if (y->m_left)y->m_left->m_parent = x;// 将 “x的父亲” 设为 “y的父亲”y->m_parent = x->m_parent;if (x->m_parent == nullptr)root = y;            // 如果 “x的父亲” 是空节点，则将y设为根节点else{if (x->m_parent->m_left == x)x->m_parent->m_left = y;    // 如果 x是它父节点的左孩子，则将y设为“x的父节点的左孩子”elsex->m_parent->m_right = y;    // 如果 x是它父节点的左孩子，则将y设为“x的父节点的左孩子”}// 将 “x” 设为 “y的左孩子”y->m_left = x;// 将 “x的父节点” 设为 “y”x->m_parent = y;
}/** 对红黑树的节点(y)进行右旋转** 右旋示意图(对节点y进行左旋)：*            py                               py*           /                                /*          y                                x*         /  \      --(右旋)-->      /  \                     #*        x   ry                           lx   y*       / \                                   / \                   #*      lx  rx                                rx  ry**/
template <class T>
void rbt_tree<T>::right_rotate(rbt_node<T>* &root, rbt_node<T>* y)
{// 设置x是当前节点的左孩子。rbt_node<T> *x = y->m_left;// 将 “x的右孩子” 设为 “y的左孩子”；// 如果"x的右孩子"不为空的话，将 “y” 设为 “x的右孩子的父亲”y->m_left = x->m_right;if (x->m_right)x->m_right->m_parent = y;// 将 “y的父亲” 设为 “x的父亲”x->m_parent = y->m_parent;if (y->m_parent == nullptr)root = x;            // 如果 “y的父亲” 是空节点，则将x设为根节点else{if (y == y->m_parent->m_right)y->m_parent->m_right = x;    // 如果 y是它父节点的右孩子，则将x设为“y的父节点的右孩子”elsey->m_parent->m_left = x;    // (y是它父节点的左孩子) 将x设为“x的父节点的左孩子”}// 将 “y” 设为 “x的右孩子”x->m_right = y;// 将 “y的父节点” 设为 “x”y->m_parent = x;
}/** 红黑树插入修正函数** 在向红黑树中插入节点之后(失去平衡)，再调用该函数；* 目的是将它重新塑造成一颗红黑树。** 参数说明：*     root 红黑树的根*     node 插入的结点        // 对应《算法导论》中的z*/
template <class T>
void rbt_tree<T>::insert_fix_up(rbt_node<T>* &root, rbt_node<T>* node)
{/* 父节点和祖父节点 */rbt_node<T> *parent = nullptr, *gparent = nullptr;// 若“父节点存在，并且父节点的颜色是红色”// 红黑树的相邻节点不能都是红色的,所以要进行修复了while ((parent = rb_parent(node)) && rb_is_red(parent)){// 先获取祖父节点gparent = rb_parent(parent);//若“父节点”是“祖父节点的左孩子”if (parent == gparent->m_left){// Case 1条件：叔叔节点是红色// 将父节点和叔叔节点全都设置为黑色,设置祖父节点为红色,讲当前节点指针指向祖父节点,进行上面的修复操作{// 获取叔叔节点rbt_node<T> *uncle = gparent->m_right;// 存在叔叔节点而且叔叔节点还是红色if (uncle && rb_is_red(uncle)){// 父节点和叔叔节点设置为黑色rb_set_black(uncle);rb_set_black(parent);// 祖父节点设置为红色rb_set_red(gparent);// 转到祖父节点进行修复node = gparent;continue;}}// Case 2条件：叔叔是黑色，且当前节点是右孩子if (parent->m_right == node){// 进行左旋转,变化为Case 3rbt_node<T> *tmp = nullptr;left_rotate(root, parent);// 节点互换了tmp = parent;parent = node;node = tmp;}// Case 3条件：叔叔是黑色，且当前节点是左孩子。// 设置父节点为黑色,将祖父节点设置为红色,再将祖父节点进行右旋转rb_set_black(parent);rb_set_red(gparent);right_rotate(root, gparent);}else//若“z的父节点”是“z的祖父节点的右孩子”{// Case 1条件：叔叔节点是红色{// 获取叔叔节点rbt_node<T> *uncle = gparent->m_left;// 存在叔叔节点且为红色的if (uncle && rb_is_red(uncle)){// 将叔叔节点和父节点设置为黑色rb_set_black(uncle);rb_set_black(parent);// 将祖父节点设置为红色rb_set_red(gparent);// 设置祖父节点为当前节点,继续修复node = gparent;continue;}}// Case 2条件：叔叔是黑色，且当前节点是左孩子// 对父节点进行右旋转,将父节点设置为当前节点if (parent->m_left == node){rbt_node<T> *tmp;right_rotate(root, parent);tmp = parent;parent = node;node = tmp;}// Case 3条件：叔叔是黑色，且当前节点是右孩子。// 将父节点设置为黑色rb_set_black(parent);// 将祖父节点设置为红色rb_set_red(gparent);// 左旋操作left_rotate(root, gparent);}}// 将根节点设为黑色rb_set_black(root);
}/** 将结点插入到红黑树中** 参数说明：*     root 红黑树的根结点*     node 插入的结点        // 对应《算法导论》中的node*/
template <class T>
void rbt_tree<T>::insert(rbt_node<T>* &root, rbt_node<T>* node)
{rbt_node<T> *y = nullptr;rbt_node<T> *x = root;// 1. 将红黑树当作一颗二叉查找树，将节点添加到二叉查找树中。// 从根节点开始查找合适的插入位置while (x){// 保存上一个指针y = x;if (node->m_key < x->m_key)	// 小于放左边x = x->m_left;else                                              // 大于放右边x = x->m_right;}// 设置父节点node->m_parent = y;// 放置在哪一边if (y){if (node->m_key < y->m_key)	// 小于放左边y->m_left = node;else                                             // 大于放右边y->m_right = node;}else// 直接设置为根节点root = node;// 2. 设置节点的颜色为红色node->m_color = rbt_red;// 3. 将它重新修正为一颗二叉查找树insert_fix_up(root, node);
}/** 将结点(key为节点键值)插入到红黑树中** 参数说明：*     tree 红黑树的根结点*     key 插入结点的键值*/
template <class T>
void rbt_tree<T>::insert(T key)
{/* 申请内存空间 */rbt_node<T> *z = new rbt_node<T>(key, rbt_black, nullptr, nullptr, nullptr);error(z, "insert new error");/* 插入节点 */insert(m_root, z);
}/** 红黑树删除修正函数** 在从红黑树中删除插入节点之后(红黑树失去平衡)，再调用该函数；* 目的是将它重新塑造成一颗红黑树。** 参数说明：*     root 红黑树的根*     node 待修正的节点*/
template <class T>
void rbt_tree<T>::remove_fix_up(rbt_node<T>* &root, rbt_node<T> *node, rbt_node<T> *parent)
{rbt_node<T> *other = nullptr;while ((!node || rb_is_black(node)) && node != root){if (parent->m_left == node){other = parent->m_right;if (rb_is_red(other)){// Case 1: x的兄弟w是红色的rb_set_black(other);rb_set_red(parent);left_rotate(root, parent);other = parent->m_right;}if ((!other->m_left || rb_is_black(other->m_left)) &&(!other->m_right || rb_is_black(other->m_right))){// Case 2: x的兄弟w是黑色，且w的俩个孩子也都是黑色的rb_set_red(other);node = parent;parent = rb_parent(node);}else{if (!other->m_right || rb_is_black(other->m_right)){// Case 3: x的兄弟w是黑色的，并且w的左孩子是红色，右孩子为黑色。rb_set_black(other->m_left);rb_set_red(other);right_rotate(root, other);other = parent->m_right;}// Case 4: x的兄弟w是黑色的；并且w的右孩子是红色的，左孩子任意颜色。rb_set_color(other, rb_color(parent));rb_set_black(parent);rb_set_black(other->m_right);left_rotate(root, parent);node = root;break;}}else{other = parent->m_left;if (rb_is_red(other)){// Case 1: x的兄弟w是红色的rb_set_black(other);rb_set_red(parent);right_rotate(root, parent);other = parent->m_left;}if ((!other->m_left || rb_is_black(other->m_left)) &&(!other->m_right || rb_is_black(other->m_right))){// Case 2: x的兄弟w是黑色，且w的俩个孩子也都是黑色的rb_set_red(other);node = parent;parent = rb_parent(node);}else{if (!other->m_left || rb_is_black(other->m_left)){// Case 3: x的兄弟w是黑色的，并且w的左孩子是红色，右孩子为黑色。rb_set_black(other->m_right);rb_set_red(other);left_rotate(root, other);other = parent->m_left;}// Case 4: x的兄弟w是黑色的；并且w的右孩子是红色的，左孩子任意颜色。rb_set_color(other, rb_color(parent));rb_set_black(parent);rb_set_black(other->m_left);right_rotate(root, parent);node = root;break;}}}if (node)rb_set_black(node);
}/** 删除结点(node)，并返回被删除的结点** 参数说明：*     root 红黑树的根结点*     node 删除的结点*/
template <class T>
void rbt_tree<T>::remove(rbt_node<T>* &root, rbt_node<T> *node)
{rbt_node<T> *child = nullptr, *parent = nullptr;rbt_color color = rbt_red;// 被删除节点的"左右孩子都不为空"的情况。if (node->m_left && node->m_right){// 被删节点的后继节点。(称为"取代节点")// 用它来取代"被删节点"的位置，然后再将"被删节点"去掉。rbt_node<T> *replace = node;// 获取后继节点,查找删除节点的右节点的最左节点replace = replace->m_right;while (replace->m_left != NULL)replace = replace->m_left;// "node节点"不是根节点(只有根节点不存在父节点)// 用替换节点覆盖删除节点if (rb_parent(node)){if (rb_parent(node)->left == node)		// 删除节点是左节点rb_parent(node)->left = replace;else                                                     //	删除节点是右节点rb_parent(node)->right = replace;}else // "node节点"是根节点，更新根节点。root = replace;// child是"取代节点"的右孩子，也是需要"调整的节点"。// "取代节点"肯定不存在左孩子！因为它是一个后继节点。child = replace->m_right;parent = rb_parent(replace);// 保存"取代节点"的颜色color = rb_color(replace);// "被删除节点"是"它的后继节点的父节点"if (parent == node)parent = replace;else{// child不为空if (child)rb_set_parent(child, parent);parent->m_left = child;replace->m_right = node->m_right;rb_set_parent(node->m_right, replace);}replace->m_parent = node->m_parent;replace->m_color = node->m_color;replace->m_left = node->m_left;node->m_left->parent = replace;if (color == rbt_black)remove_fix_up(root, child, parent);delete node;node = nullptr;return;}if (node->m_left)child = node->m_left;elsechild = node->m_right;// 保存父节点parent = node->m_parent;// 保存"取代节点"的颜色color = node->m_color;// 设置父节点if (child)child->m_parent = parent;// "node节点"不是根节点if (parent){if (parent->m_left == node)parent->m_left = child;elseparent->m_right = child;}elseroot = child;// 黑色就需要修复if (color == rbt_black)remove_fix_up(root, child, parent);delete node;node = nullptr;return;
}/** 删除红黑树中键值为key的节点** 参数说明：*     tree 红黑树的根结点*/
template <class T>
void rbt_tree<T>::remove(T key)
{// 查找key对应的节点(node)，找到的话就删除该节点rbt_node<T> *node = search(m_root, key);if (node)remove(m_root, node);
}/** 销毁红黑树*/
template <class T>
void rbt_tree<T>::destroy(rbt_node<T>* &tree)
{if (tree == nullptr) return;if (tree->m_left)		// 销毁左节点return destroy(tree->m_left);if (tree->m_right)	// 销毁右节点return destroy(tree->m_right);// 释放并清空delete tree;tree = nullptr;
}template <class T>
void rbt_tree<T>::destroy()
{destroy(m_root);
}/** 打印"二叉查找树"** key        -- 节点的键值* direction  --  0，表示该节点是根节点;*               -1，表示该节点是它的父结点的左孩子;*                1，表示该节点是它的父结点的右孩子。*/
template <class T>
void rbt_tree<T>::print(rbt_node<T>* tree, T key, int direction)
{if (tree){if (direction == 0)    // tree是根节点cout << setw(2) << tree->m_key << "(B) is root" << endl;else							 // tree是分支节点cout << setw(2) << tree->m_key << (rb_is_red(tree) ? "(R)" : "(B)") << " is " << setw(2) << key << "'s " << setw(12) << (direction == 1 ? "right child" : "left child") << endl;print(tree->m_left, tree->m_key, -1);print(tree->m_right, tree->m_key, 1);}
}template <class T>
void rbt_tree<T>::print()
{if (m_root)print(m_root, m_root->m_key, 0);
}