371 lines
9.0 KiB
C++
371 lines
9.0 KiB
C++
#pragma once
|
||
#include <iostream>
|
||
#include "TreeNode.h"
|
||
template <class T>
|
||
class RBTree {
|
||
private:
|
||
TreeNode<T>* root;
|
||
TreeNode<T>* nul;
|
||
int size;
|
||
|
||
void initNul() {
|
||
nul = new TreeNode<T>(0, nul, nul, nul);
|
||
nul->color = Color::RED;
|
||
}
|
||
void inorder(TreeNode<T>* x);
|
||
void dispose(TreeNode<T>* x) {
|
||
if (x == nullptr || x == nul) return;
|
||
dispose(x->left);
|
||
dispose(x->right);
|
||
delete x;
|
||
//重点不能delete nul（作为全局的哨兵）
|
||
}
|
||
public:
|
||
RBTree() : size(0) { initNul(); root = nul; }
|
||
RBTree(TreeNode<T>* r, int s) : root(r), size(s) { initNul(); }
|
||
~RBTree() {
|
||
dispose(root);
|
||
root = nullptr;
|
||
delete nul;
|
||
//删除哨兵，需要删除
|
||
}
|
||
|
||
TreeNode<T>* RotateLeft(TreeNode<T>* x);//RR
|
||
TreeNode<T>* RotateRight(TreeNode<T>* x);//LL
|
||
|
||
|
||
void insert(const T& e);
|
||
void fixinsert(TreeNode<T>* x);
|
||
void fixdelete(TreeNode<T>* x);
|
||
void inorder() { inorder(root); }
|
||
void erase(const T& e);
|
||
TreeNode<T>* find(const T& e);
|
||
};
|
||
|
||
template<class T>
|
||
inline TreeNode<T>* RBTree<T>::RotateLeft(TreeNode<T>* x) {
|
||
TreeNode<T>* y = x->right; // 右孩子
|
||
TreeNode<T>* T2 = y->left; // y 的左子树暂存
|
||
|
||
// 左旋
|
||
y->left = x;
|
||
x->right = T2;
|
||
|
||
// 更新 parent 指针
|
||
if (T2) T2->parent = x; // 如果 T2 存在
|
||
y->parent = x->parent; // y 接上 x 的父节点
|
||
|
||
if (x->parent == nul) { // x 是根节点
|
||
root = y;
|
||
}
|
||
else if (x == x->parent->left) {
|
||
x->parent->left = y;
|
||
}
|
||
else {
|
||
x->parent->right = y;
|
||
}
|
||
|
||
x->parent = y;
|
||
|
||
return y; // 返回新子树根
|
||
}
|
||
|
||
|
||
template<class T>
|
||
TreeNode<T>* RBTree<T>::RotateRight(TreeNode<T>* x) {
|
||
//有parent，不需要递归，直接修改parent
|
||
TreeNode<T>* y = x->left;
|
||
TreeNode<T>* T2 = y->right;
|
||
|
||
// 旋转
|
||
y->right = x;
|
||
x->left = T2;
|
||
|
||
// 更新 parent 指针
|
||
if (T2) T2->parent = x; // 如果 T2 存在
|
||
y->parent = x->parent; // y 接上 x 的父节点
|
||
|
||
if (x->parent == nul) { // x 是根节点
|
||
root = y;
|
||
}
|
||
else if (x == x->parent->left) {
|
||
x->parent->left = y;
|
||
}
|
||
else {
|
||
x->parent->right = y;
|
||
}
|
||
x->parent = y;
|
||
return y; // 返回新子树根
|
||
}
|
||
|
||
|
||
template<class T>
|
||
inline void RBTree<T>::inorder(TreeNode<T>* x)
|
||
{
|
||
if (x == nul) return;
|
||
inorder(x->left);
|
||
std::cout << x->element << " ";
|
||
inorder(x->right);
|
||
}
|
||
|
||
template<class T>
|
||
inline void RBTree<T>::insert(const T& e)
|
||
{
|
||
TreeNode<T>* p = new TreeNode<T>(e, nul, nul, nul);
|
||
|
||
TreeNode<T>* current = root;
|
||
TreeNode<T>* parent = nul;
|
||
//BST插入，有parent就不需要递归节点了
|
||
while (current != nul) {
|
||
parent = current;
|
||
if (e < parent->element) current = parent->left;
|
||
else current = parent->right;
|
||
}
|
||
|
||
p->parent = parent;
|
||
if (parent == nul) root = p;
|
||
if (parent == nul) root = p;
|
||
else if (e < parent->element) parent->left = p;
|
||
else parent->right = p;
|
||
|
||
//FixInsert
|
||
fixinsert(p);
|
||
}
|
||
|
||
template<class T>
|
||
inline void RBTree<T>::fixinsert(TreeNode<T>* x)
|
||
{
|
||
while (x != root && x->parent->color == Color::RED) {
|
||
//看叔叔
|
||
if (x->parent->parent->left == x->parent) {
|
||
//左节点
|
||
if (x->parent->parent->right->color == Color::RED) {
|
||
//叔叔红
|
||
x->parent->color = x->parent->parent->right->color = Color::BLACK;
|
||
x->parent->parent->color = Color::RED;
|
||
//fixinsert(x->parent->parent);递归不好，容易重复判断，不易处理根节点颜色
|
||
x = x->parent->parent;
|
||
continue;
|
||
}
|
||
if (x == x->parent->left) {
|
||
//LL//先变色祖父，父亲
|
||
x->parent->parent->color = Color::RED;
|
||
x->parent->color = Color::BLACK;
|
||
//再右旋
|
||
RotateRight(x->parent->parent);
|
||
}
|
||
//LR
|
||
if (x == x->parent->right) {
|
||
//变成LL
|
||
x = x->parent;
|
||
RotateLeft(x);
|
||
//不需要再LL了，下一个循环自动变成LL
|
||
}
|
||
}
|
||
//else对称情况，父亲为右节点
|
||
else {
|
||
//1-叔叔红
|
||
if (x->parent->parent->left->color == Color::RED) {
|
||
x->parent->parent->left->color = Color::BLACK;
|
||
x->parent->color = Color::BLACK;
|
||
x->parent->parent->color = Color::RED;
|
||
x = x->parent->parent;
|
||
}
|
||
//2-叔叔黑
|
||
else {
|
||
//2-1-RL -> 调整为RR，再做下一步操作
|
||
if (x == x->parent->left) {
|
||
x = x->parent;
|
||
RotateRight(x);
|
||
}
|
||
|
||
//3-RR，接着上面的2-1
|
||
x->parent->color = Color::BLACK;
|
||
x->parent->parent->color = Color::RED;
|
||
RotateLeft(x->parent->parent);
|
||
}
|
||
}
|
||
}
|
||
root->color = Color::BLACK;
|
||
}
|
||
// x 表示被删除节点的替代节点（可能是 nul 哨兵）
|
||
// 只有当被删节点或替代节点是黑色时，才需要进入 fixdelete
|
||
template<class T>
|
||
void RBTree<T>::fixdelete(TreeNode<T>* x) {
|
||
while (x != root && x->color == Color::BLACK) {
|
||
// 情况1：x 是左孩子
|
||
if (x == x->parent->left) {
|
||
TreeNode<T>* w = x->parent->right; // x 的兄弟节点
|
||
|
||
// case 1: 兄弟是红色
|
||
// -> 把兄弟染黑，父亲染红，然后左旋
|
||
// -> 转化为兄弟是黑色的情况
|
||
if (w->color == Color::RED) {
|
||
w->color = Color::BLACK;
|
||
x->parent->color = Color::RED;
|
||
RotateLeft(x->parent);
|
||
w = x->parent->right; // 更新兄弟
|
||
}
|
||
|
||
// case 2: 兄弟是黑色，并且两个侄子（w->left, w->right）都是黑色
|
||
// -> 把兄弟染红，相当于把 "多余的黑色" 往上传递
|
||
// -> 修复过程继续向上进行
|
||
if (w->left->color == Color::BLACK && w->right->color == Color::BLACK) {
|
||
w->color = Color::RED;
|
||
x = x->parent; // 往上修复
|
||
}
|
||
else {
|
||
// case 3: 兄弟是黑色，且兄弟的右孩子是黑色，左孩子是红色
|
||
// -> 先把兄弟染红，左孩子染黑，然后右旋兄弟
|
||
// -> 转化为 case 4
|
||
if (w->right->color == Color::BLACK) {
|
||
w->left->color = Color::BLACK;
|
||
w->color = Color::RED;
|
||
RotateRight(w);
|
||
w = x->parent->right; // 更新兄弟
|
||
}
|
||
|
||
// case 4: 兄弟是黑色，且兄弟的右孩子是红色
|
||
// -> 用兄弟替换父亲颜色，父亲染黑，兄弟的右孩子染黑
|
||
// -> 左旋父亲，修复完成，直接退出循环
|
||
w->color = x->parent->color;
|
||
x->parent->color = Color::BLACK;
|
||
w->right->color = Color::BLACK;
|
||
RotateLeft(x->parent);
|
||
x = root; // 结束循环
|
||
}
|
||
}
|
||
// 情况2：x 是右孩子（对称情况）
|
||
else {
|
||
TreeNode<T>* w = x->parent->left;
|
||
|
||
if (w->color == Color::RED) {
|
||
w->color = Color::BLACK;
|
||
x->parent->color = Color::RED;
|
||
RotateRight(x->parent);
|
||
w = x->parent->left;
|
||
}
|
||
|
||
if (w->right->color == Color::BLACK && w->left->color == Color::BLACK) {
|
||
w->color = Color::RED;
|
||
x = x->parent;
|
||
}
|
||
else {
|
||
if (w->left->color == Color::BLACK) {
|
||
w->right->color = Color::BLACK;
|
||
w->color = Color::RED;
|
||
RotateLeft(w);
|
||
w = x->parent->left;
|
||
}
|
||
|
||
w->color = x->parent->color;
|
||
x->parent->color = Color::BLACK;
|
||
w->left->color = Color::BLACK;
|
||
RotateRight(x->parent);
|
||
x = root;
|
||
}
|
||
}
|
||
}
|
||
|
||
// 最后保证 x 是黑色（可能是根或者 nul）
|
||
x->color = Color::BLACK;
|
||
}
|
||
|
||
template<class T>
|
||
inline void RBTree<T>::erase(const T& e)
|
||
{
|
||
TreeNode<T>* z = find(e);
|
||
//1-BST删除重置，找本节点右树最小或者左树最大
|
||
TreeNode<T>* y = z; // 实际要删除的节点（可能是 z，也可能是 z 的后继）
|
||
TreeNode<T>* x; // y 的替代节点（可能是子节点，也可能是 nul）
|
||
Color yOriginalColor = y->color; // 记录 y 的原始颜色，决定是否需要 fixdelete
|
||
// case 1: z 没有左子树 -> 用右子树替代 z
|
||
if (z->left == nul) {
|
||
x = z->right; // 替代节点
|
||
// 把 z 从父亲断开
|
||
if (z->parent == nul) root = x;
|
||
else if (z == z->parent->left) z->parent->left = x;
|
||
else z->parent->right = x;
|
||
x->parent = z->parent;
|
||
}
|
||
// case 2: z 没有右子树 -> 用左子树替代 z
|
||
else if (z->right == nul) {
|
||
x = z->left;
|
||
if (z->parent == nul) root = x;
|
||
else if (z == z->parent->left) z->parent->left = x;
|
||
else z->parent->right = x;
|
||
x->parent = z->parent;
|
||
}
|
||
// case 3: z 左右孩子都有 -> 找到 z 的中序后继
|
||
else {
|
||
// 找 z 的后继（右子树最左节点）
|
||
y = z->right;
|
||
while (y->left != nul) y = y->left;
|
||
yOriginalColor = y->color; // 记录后继的颜色
|
||
x = (y->right != nul ? y->right : nul); // 后继的右孩子作为替代
|
||
|
||
// 如果 y 不是 z 的直接右孩子
|
||
if (y->parent != z) {
|
||
// 把 y->right 替代 y
|
||
if (y->parent->left == y) y->parent->left = x;
|
||
else y->parent->right = x;
|
||
x->parent = y->parent;
|
||
|
||
// 把 z->right 挂到 y 的右子树
|
||
y->right = z->right;
|
||
y->right->parent = y;
|
||
}
|
||
|
||
// 把 z->left 挂到 y 的左子树
|
||
if (z->parent == nul) root = y;
|
||
else if (z == z->parent->left) z->parent->left = y;
|
||
else z->parent->right = y;
|
||
y->parent = z->parent;
|
||
|
||
y->left = z->left;
|
||
y->left->parent = y;
|
||
y->color = z->color;
|
||
}
|
||
delete z;
|
||
|
||
//2-RBT性质重置
|
||
// 如果删除的实际节点是黑色，可能破坏红黑树性质 → fixdelete
|
||
if (yOriginalColor == Color::BLACK) {
|
||
fixdelete(x);
|
||
}
|
||
|
||
}
|
||
template<class T>
|
||
inline TreeNode<T>* RBTree<T>::find(const T& e)
|
||
{
|
||
TreeNode<T>* p = root;
|
||
while (p != nul) {
|
||
if (e < p->element) p = p->left;
|
||
else if (e > p->element) p = p->right;
|
||
else break;
|
||
}
|
||
if (!p) p = nul;
|
||
return p;
|
||
}
|
||
/*
|
||
推荐写法
|
||
if (x->parent == x->parent->parent->left) { // 父在左
|
||
TreeNode<T>* y = x->parent->parent->right; // 叔叔
|
||
if (y->color == Color::RED) { // Case 1: 叔红
|
||
x->parent->color = Color::BLACK;
|
||
y->color = Color::BLACK;
|
||
x->parent->parent->color = Color::RED;
|
||
x = x->parent->parent; // 上移祖父
|
||
} else {
|
||
if (x == x->parent->right) { // Case 2: LR
|
||
x = x->parent;
|
||
RotateLeft(x);
|
||
}
|
||
// Case 3: LL
|
||
x->parent->color = Color::BLACK;
|
||
x->parent->parent->color = Color::RED;
|
||
RotateRight(x->parent->parent);
|
||
}
|
||
按照顺序写，这里可以看else（父亲为右节点）的思路
|
||
*/
|