红黑树

/*
* 红黑树
*      1 提供基本的操作集：
*          得到最大结点
*          得到最小结点
*          得到后继结点
*          得到前驱结点
*          插入结点
*          删除结点
*      2 可以扩展的操作
*          输出一个有序列
* 本例中没有对重复的元素进行控制
*
* NIL 这个思想是从《算法导论》中来的
* 有了 NIL 所有的结点都是内结点，便于算法的操作
*
* Successor 算法需要明白两种情况，没有右孩子，有右孩子
* InsertNode 算法只插入叶子
* DeleteNode 算法不删除既有右孩子又有左孩子的结点，这种情况需要转化
*
*/

#include <cstdio>

#define RED    1
#define BLACK  0
struct RBNode {
int key;
int color;
RBNode *pParent;
RBNode *pLChild;
RBNode *pRChild;
};

class CRBTree {
public:
CRBTree() {
NIL = new RBNode;
NIL->color = BLACK;
NIL->pParent = NIL;
NIL->pLChild = NIL;
NIL->pRChild = NIL;

root = NIL;
}

~CRBTree() {
delete NIL;
}

int GetMaxKey() {
RBNode *node = GetMaxNode(root);
if (node == NIL) return 0;
else return node->key;
}

int GetMinKey() {
RBNode *node = GetMinNode(root);
if (node == NIL) return 0;
else return node->key;
}

void InsertKey(int key) {
RBNode *node = new RBNode;
node->key = key;
node->color = RED;
node->pParent = NIL;
node->pLChild = NIL;
node->pRChild = NIL;

InsertNode(node);
}

void DeleteKey(int key) {
RBNode *node;
while ((node=FindInTree(key)) != NIL) {
node = DeleteNode(node);
if (node != NIL) delete node;
}
}

bool KeyInTree(int key) {
return FindInTree(key) != NIL;
}

private:
RBNode *GetMaxNode(RBNode *node) {
RBNode *p = node;
while (p->pRChild != NIL) {
p = p->pRChild;
}
return p;
}

RBNode *GetMinNode(RBNode *node) {
RBNode *p = node;
while (p->pLChild != NIL) {
p = p->pLChild;
}
return p;
}

RBNode *FindInTree(int key) {
RBNode *p = root;
while (p!=NIL && p->key!=key) {
if (p->key < key) p = p->pRChild;
else p = p->pLChild;
}
return  p;
}

RBNode *Successor(RBNode *node) {
if (node->pRChild != NIL) return GetMinNode(node->pRChild);

RBNode *p = node;
RBNode *pParent = p->pParent;
while (pParent!=NIL && pParent->pRChild==p) {
p = pParent;
pParent = p->pParent;
}

return pParent;
}

RBNode *Predecessor(RBNode *node) {
if (node->pLChild != NIL) return GetMaxNode(node->pLChild);

RBNode *p = node;
RBNode *pParent = p->pParent;
while (pParent!=NIL && pParent->pLChild==p) {
p = pParent;
pParent = p->pParent;
}
return pParent;
}

void LeftRotate(RBNode *node) {
// x y 都不是 NIL
RBNode *x = node;
RBNode *y = node->pRChild;

y->pParent = x->pParent;
if (x->pParent == NIL) root = y;
else if (x->pParent->pLChild == x) x->pParent->pLChild = y;
else x->pParent->pRChild = y;

x->pRChild = y->pLChild;
if (x->pRChild != NIL) x->pRChild->pParent = x;

y->pLChild = x;
x->pParent = y;
}

void RightRotate(RBNode *node) {
// x y 都不是 NIL
RBNode *x = node;
RBNode *y = node->pLChild;
y->pParent = x->pParent;
if (x->pParent == NIL) root = y;
else if (x->pParent->pLChild == x) x->pParent->pLChild = y;
else x->pParent->pRChild = y;

x->pLChild = y->pRChild;
if (x->pLChild != NIL) x->pLChild->pParent = x;

y->pRChild = x;
x->pParent = y;
}

void InsertNode(RBNode *node) {
RBNode *p = root;
RBNode *pParent = NIL;

while (p != NIL) {
pParent = p;
if (p->key < node->key) p = p->pRChild;
else p = p->pLChild;
}

node->pParent = pParent;
if (pParent == NIL) root = node;
else if (pParent->key < node->key) pParent->pRChild = node;
else pParent->pLChild = node;

node->pLChild = NIL;
node->pRChild = NIL;
node->color = RED;

InsertFixup(node);
}

void InsertFixup(RBNode *node) {
RBNode *z = node;
while (z->pParent->color==RED) {
if (z->pParent==z->pParent->pParent->pLChild) {
if (z->pParent->pParent->pRChild->color==RED) {
z->pParent->color =  BLACK;
z->pParent->pParent->pRChild->color= BLACK;
z->pParent->pParent->color = RED;
z = z->pParent->pParent;
}
else {
if (z->pParent->pRChild == z) {
z = z->pParent;
LeftRotate(z);
}

z->pParent->color = BLACK;
z->pParent->pParent->color = RED;
RightRotate(z->pParent->pParent);
}
}
else {
if (z->pParent->pParent->pLChild->color == RED) {
z->pParent->color = BLACK;
z->pParent->pParent->pLChild->color = BLACK;
z->pParent->pParent->color = RED;
z = z->pParent->pParent;
}
else {
if (z->pParent->pLChild == z) {
z = z->pParent;
RightRotate(z);
}

z->pParent->color = BLACK;
z->pParent->pParent->color = RED;
LeftRotate(z->pParent->pParent);
}
}
}

root->color = BLACK;
}

RBNode *DeleteNode(RBNode *node) {
RBNode *p = NIL;
if (node->pLChild==NIL || node->pRChild==NIL) p = node;
else p = Successor(node);

RBNode *pChild = NIL;
if (p->pLChild != NIL)  pChild= p->pLChild;
else pChild = p->pRChild;

pChild->pParent = p->pParent;
if (p->pParent == NIL) root = NIL;
else if (p->pParent->pLChild == p) p->pParent->pLChild = pChild;
else p->pParent->pRChild = pChild;

if (p != node) {
int tmp = node->key;
node->key = p->key;
p->key = tmp;
}

if (p->color == BLACK) DeleteFixup(pChild);

if (pChild == NIL) pChild->pParent = NIL;

return p;
}

void DeleteFixup(RBNode *x) {
RBNode *w = NIL;
while (x != root && x->color==BLACK) {
if (x->pParent->pLChild == x) {
w = x->pParent->pRChild;
if (w->color == RED) {
w->color = BLACK;
x->pParent->color = RED;
LeftRotate(x->pParent);
w = x->pParent->pRChild;
}

if (w->pLChild->color==BLACK && w->pRChild->color==BLACK) {
w->color = RED;
x = x->pParent;
}
else {
if (w->pRChild->color == BLACK) {
w->pLChild->color = BLACK;
w->color = RED;
RightRotate(w);
w = w->pParent;
}

w->color = x->pParent->color;
x->pParent->color = BLACK;
w->pRChild->color = BLACK;
LeftRotate(x->pParent);
x = root;
}

}
else {
w = x->pParent->pLChild;
if (w->color == RED) {
w->color = BLACK;
x->pParent->color = RED;
RightRotate(x->pParent);
w = x->pParent->pLChild;
}

if (w->pLChild->color==BLACK && w->pRChild->color==BLACK) {
w->color = RED;
x = x->pParent;
}
else {
if (w->pLChild == BLACK) {
w->pRChild->color = BLACK;
w->color = RED;
LeftRotate(w);
w = w->pParent;
}

w->color = x->pParent->color;
x->pParent->color = BLACK;
w->pLChild->color = BLACK;
RightRotate(x->pParent);
x = root;
}
}
}
x->color = BLACK;
}

private:
RBNode *NIL;
RBNode *root;
};

int main() {

int a[10] = {2, 5, 7, 8, 1, 9 ,10, 3, 4, 6};

CRBTree btree;
for (int i=0; i<10; ++i)
btree.InsertKey(a[i]);

printf("%d %d\n", btree.GetMinKey(), btree.GetMaxKey());

btree.DeleteKey(1);
btree.DeleteKey(10);
printf("%d %d\n", btree.GetMinKey(), btree.GetMaxKey());

btree.DeleteKey(5);
if (btree.KeyInTree(5))
puts("have found 5 in btree");
else
puts("can not find 5 in btree");

if (btree.KeyInTree(4))
puts("have found 4 in btree");
else
puts("can not find 4 in btree");

return 0;
}