平衡树之RB-tree

#include <memory>

template<class T>
struct rb_node
{
T key;
bool color;//true red | false black
std::shared_ptr<rb_node> lchild, rchild, parent;

rb_node(T key, bool color = true, std::shared_ptr<rb_node> lchild = nullptr,
std::shared_ptr<rb_node> rchild = nullptr, std::shared_ptr<rb_node> parent = nullptr)
:key(key)//此处不能使用this：现在对象还没被构建起来
{
//赋值，初始化在初始化列表中执行
this->color = color;
this->lchild = lchild;
this->rchild = rchild;
this->parent = parent;
}
};

template<class T>
class rb_tree
{
private:
std::shared_ptr<rb_node<T>> root;
std::shared_ptr<rb_node<T>> nil;//叶节点
void left_rotation(std::shared_ptr<rb_node<T>> a)
{
if (a == nullptr) return;
std::shared_ptr<rb_node<T>> b = a->rchild;
if (b == nullptr) return;
if (b->lchild != nullptr)
b->lchild->parent = a;
a->rchild = b->lchild;

if (a->parent == nil)
root = b;//rbtree的root以nil为父节点
else
{
if (a->parent->lchild == a)
a->parent->lchild = b;
else
a->parent->rchild = b;
}
b->parent = a->parent;

b->lchild = a;
a->parent = b;
}
void right_rotation(std::shared_ptr<rb_node<T>> a)
{
if (a == nullptr) return;
std::shared_ptr<rb_node<T>> b = a->lchild;
if (b == nullptr) return;
if (b->rchild != nullptr)
b->rchild->parent = a;
a->lchild = b->rchild;

if (a->parent == nil)
root = b;
else
{
if (a->parent->lchild == a)
a->parent->lchild = b;
else
a->parent->rchild = b;
}
b->parent = a->parent;

b->rchild = a;
a->parent = b;
}

public:
rb_tree()
{
root = nullptr;
T key;
nil = std::make_shared<rb_node<T>>(key, false);
}
void insert(T key)
{
std::shared_ptr<rb_node<T>> tmp = std::make_shared<rb_node<T>>(key, true, nil, nil, nil);
std::shared_ptr<rb_node<T>> ptr = root;

//情况1：树为空
if (ptr == nullptr)
{
tmp->color = false;
root = tmp;
return;
}
while (true)
{
if (key <= ptr->key)
{
if (ptr->lchild == nil) break;
ptr = ptr->lchild;
}
else
{
if (ptr->rchild == nil) break;
ptr = ptr->rchild;
}
}

if (key <= ptr->key)
ptr->lchild = tmp;
else
ptr->rchild = tmp;
tmp->parent = ptr;

while(true)
{
if (ptr == nil)//注意root可能被情况三修改为red，记得加特判
{
tmp->color = false;
root = tmp;
return;
}

//情况2：插入节点的父节点为黑色
if (!ptr->color) return;

/*情况3：插入节点的父节点和叔节点都存在且都为红色*/
if (ptr->parent->lchild->color && ptr->parent->rchild->color)
{
ptr->parent->color = true;
ptr->parent->lchild->color = false;
ptr->parent->rchild->color = false;

tmp = ptr->parent;
ptr = tmp->parent;
continue;
}
if (ptr->parent->lchild == ptr)
{
//情况4：右旋
if (tmp == ptr->lchild)
{
ptr->parent->color = true;
ptr->color = false;
right_rotation(ptr->parent);
return;
}
else
{
//情况5：左旋 + 右旋
left_rotation(ptr);
tmp = ptr;
ptr = tmp->parent;
continue;
}
}
else
{
//情况4：左旋
if (tmp == ptr->rchild)
{
ptr->parent->color = true;
ptr->color = false;
left_rotation(ptr->parent);
return;
}
else
{
//情况5：右旋+左旋
right_rotation(ptr);
tmp = ptr;
ptr = tmp->parent;
continue;
}
}
}
}
void earse(T key)
{
//待续
}
};