《算法导论》中红黑树的C语言实现

《算法导论》好是好，只是看伪代码太头痛了，总想要是有个C语言版的《算法导论》就好了。

前几天研究了红黑树，把它翻译成C语言。在VC6.0上编译测试通过。并且写了一个CView类来显示效果。

头文件 rb_tree.h

#ifndef RB_TREE_H
#define RB_TREE_H

enum NODECOLOR {RED = 0, BLACK = 1};

typedef int NODEDATA;

struct rb_tree_node
{
NODECOLOR color;
int key;
NODEDATA info;
rb_tree_node* left;
rb_tree_node* right;
rb_tree_node* parent;
};

struct rb_tree
{
rb_tree_node* root;
rb_tree_node* nil;
};

void init_rb_tree(rb_tree* t);

void left_rotate(rb_tree* t, rb_tree_node* x);
void right_rotate(rb_tree* t, rb_tree_node* x);

void rb_insert(rb_tree* t, int key, NODEDATA info = 0);
void rb_insert_(rb_tree* t, rb_tree_node* z);
void rb_insert_fixup(rb_tree* t, rb_tree_node* z);

void rb_delete(rb_tree* t, int key);
void rb_delete_(rb_tree* t, rb_tree_node* z);
void rb_delete_fixup(rb_tree* t, rb_tree_node* x);

int get_tree_height(rb_tree* t, rb_tree_node* x);

#endif

实现文件：rb_tree.cpp

/*
红黑树的实现，根据《算法导论》伪代码翻译
作者：胡维臣 huweichen@163.com
时间：2010.6.26
*/
#include <stdlib.h>
//#include <malloc.h>

#include "rb_tree.h"

void init_rb_tree(rb_tree* t)
{
rb_tree_node* node = new rb_tree_node;

node->color = BLACK;
node->info = 0;
node->key = 0;
node->left = 0;
node->parent = 0;
node->right = 0;

t->nil = node;
t->root = t->nil;
}

int get_tree_height(rb_tree* t, rb_tree_node* x)
{
int lh = 0;
int rh = 0;

if(x->left != t->nil)
{
lh = get_tree_height(t, x->left);
}

if(x->right != t->nil)
{
rh = get_tree_height(t, x->right);
}

if(lh > rh)
{
return lh + 1;
}
else
{
return rh + 1;
}
}

void left_rotate(rb_tree* t, rb_tree_node* x)
{
rb_tree_node* y = x->right;

x->right = y->left;
if(y->left != t->nil)
{
y->left->parent = x;
}

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

y->left = x;
x->parent = y;
}

void right_rotate(rb_tree* t, rb_tree_node* x)
{
rb_tree_node* y = x->left;

//先把左子节点的右子节点作为左子节点
x->left = y->right;
if(y->right != t->nil)
{
y->right->parent = x;
}

//再把原左子节点替代自己的位置
y->parent = x->parent;
if(x->parent == t->nil)
{
t->root = y;
}
else if(x == x->parent->right)
{
x->parent->right = y;
}
else
{
x->parent->left = y;
}

//最后把自己作为原左子节点的右子节点
y->right = x;
x->parent = y;
}

void rb_insert(rb_tree* t, int key, NODEDATA info)
{
rb_tree_node* node = new rb_tree_node;
node->key = key;
node->info = info;

rb_insert_(t, node);
}

void rb_insert_(rb_tree* t, rb_tree_node* z)
{
rb_tree_node* y = t->nil;
rb_tree_node* x = t->root;

//找一个合适的插入位置
while(x != t->nil)
{
y = x;
if(z->key == x->key)
{
delete z;
return;
}
else if(z->key < x->key)
{
x = x->left;
}
else
{
x = x->right;
}
}

//插入节点
z->parent = y;
if(y == t->nil)
{
t->root = z;
}
else if(z->key < y->key)
{
y->left = z;
}
else
{
y->right = z;
}

//将新节点的其它域赋值
z->left = t->nil;
z->right = t->nil;
z->color = RED;

//修改树，以满足红黑树的性质
rb_insert_fixup(t, z);
}

void rb_insert_fixup(rb_tree* t, rb_tree_node* z)
{
rb_tree_node* y = t->nil;

//循环直到父节点为黑色
while(z->parent->color == RED)
{
//如果父节点是祖父节点的左子节点
if(z->parent == z->parent->parent->left)
{
y = z->parent->parent->right;
//如果叔叔节点也是红色，则把父节点和叔叔节点都设为黑色，把祖父节点设为红色
if(y->color == RED)
{
z->parent->color = BLACK;
y->color = BLACK;
z->parent->parent->color = RED;
z = z->parent->parent;
}
//如果叔叔节点是黑色，那么把父节点变成黑色，把祖父节点变成红色，
//然后右旋转，使祖父节点成为叔叔节点
else
{
//如果是右子节点，先左旋转，这样下面的右旋转后才会对称
if(z == z->parent->right)
{
z = z->parent;
left_rotate(t, z);
}
z->parent->color = BLACK;
z->parent->parent->color = RED;
right_rotate(t, z->parent->parent);
}
}
//下面代码是和上面对称的
else
{
y = z->parent->parent->left;
if(y->color == RED)
{
z->parent->color = BLACK;
y->color = BLACK;
z->parent->parent->color = RED;
z = z->parent->parent;
}
else
{
if(z == z->parent->left)
{
z = z->parent;
right_rotate(t, z);
}
z->parent->color = BLACK;
z->parent->parent->color = RED;
left_rotate(t, z->parent->parent);
}
}
}

t->root->color = BLACK;
}

rb_tree_node* find_node(rb_tree* t, int key)
{
rb_tree_node* x = t->root;

while(x != t->nil)
{
if(key == x->key)
{
return x;
}
else if(key < x->key)
{
x = x->left;
}
else
{
x = x->right;
}
}

return t->nil;
}

void rb_delete(rb_tree* t, int key)
{
rb_tree_node* x = find_node(t, key);
if(x != t->nil)
{
rb_delete_(t, x);
}
}

rb_tree_node* tree_minimum(rb_tree* t, rb_tree_node* x)
{
while(x->left != t->nil)
{
x = x->left;
}

return x;
}

rb_tree_node* tree_successor(rb_tree* t, rb_tree_node* x)
{
rb_tree_node* y = t->nil;
if(x->right != t->nil)
{
return tree_minimum(t, x->right);
}

y = x->parent;
while(y != t->nil && x == y->right)
{
x = y;
y = y->parent;
}

return y;
}

void rb_delete_(rb_tree* t, rb_tree_node* z)
{
rb_tree_node* y = t->nil;
rb_tree_node* x = t->nil;

//只要z有一个字节点不为空，y就等于z，否则y为z的后驱
if(z->left == t->nil || z->right == t->nil)
{
y = z;
}
else
{
y = tree_successor(t, z);
}

//把x设为y的子节点，当然也有可能为nil
if(y->left != t->nil)
{
x = y->left;
}
else
{
x = y->right;
}

//从树中删除y，但暂时不释放y所占内存，因为后面还要用到
x->parent = y->parent;
if(y->parent == t->nil)
{
t->root = x;
}
else
{
if(y == y->parent->left)
{
y->parent->left = x;
}
else
{
y->parent->right = x;
}
}

//如果y不等于z，而是z的后驱，则把y的关键字和数据拷贝到z
if(y != z)
{
z->key = y->key;
z->info = y->info;
}

//如果被删的y是黑结点，树的红黑性质遭到破坏，即黑高度变了，则需要调整
if(y->color == BLACK)
{
rb_delete_fixup(t, x);
}

//这时，可以释放y所占内存了
delete y;
}

void rb_delete_fixup(rb_tree* t, rb_tree_node* x)
{
rb_tree_node* w = t->nil;

//x的颜色为只有在黑色情况下才需要调整，如果x为红色，直接把它设为黑色就可以了
while(x != t->root && x->color == BLACK)
{
if(x == x->parent->left)
{
//把w设为兄节点
w = x->parent->right;

//如果w为红色，那么把w设为黑色，把w的父节点设为红色，然后把父节点左旋转
if(w->color == RED)
{
w->color = BLACK;
x->parent->color = RED;
left_rotate(t, x->parent);
w = x->parent->right;
}

//如果w的左右子节点都为黑色，把w设为红色，把x设为x的父节点，
//这时，新x的左右子节点的黑高度就一致了，即x是一颗红黑树，但新x的兄节点的
//黑高度仍然大一个，所以进入下一次循坏，继续调整
if(w->left->color == BLACK && w->right->color == BLACK)
{
w->color = RED;
x = x->parent;
}
else
{
//如果w的左子节点为红，右子节点为黑，则先把w设为红色，左子节点设黑色，
//然后右旋转，这样做是为下面的左旋转做好准备
if(w->right->color == BLACK)
{
w->left->color = BLACK;
w->color = RED;
right_rotate(t, w);
w = x->parent->right;
}

//把w设为父节点的颜色，这样做是为了在下面的左旋转后使新旧父节点的颜色不变，
//再把w的父节点和右子节点都设为黑色，然后左旋转，
//这样左子节点和右子节点的黑高度就一致了，不用再调整了，所以把x设为root。
w->color = x->parent->color;
x->parent->color = BLACK;
w->right->color = BLACK;
left_rotate(t, x->parent);
x = t->root;
}
}
//下面代码和上面是对称的
else
{
w = x->parent->left;
if(w->color == RED)
{
w->color = BLACK;
x->parent->color = RED;
right_rotate(t, x->parent);
w = x->parent->left;
}
if(w->right->color == BLACK && w->left->color == BLACK)
{
w->color = RED;
x = x->parent;
}
else
{
if(w->left->color == BLACK)
{
w->right->color = BLACK;
w->color = RED;
left_rotate(t, w);
w = x->parent->left;
}

w->color = x->parent->color;
x->parent->color = BLACK;
w->left->color = BLACK;
right_rotate(t, x->parent);
x = t->root;
}
}
}

x->color = BLACK;
}

用来显试效果的CTestRBTreeView类

void CTestRBTreeView::OnDraw(CDC* pDC)
{
CTestRBTreeDoc* pDoc = GetDocument();
ASSERT_VALID(pDoc);

if(rb_tree1.root == rb_tree1.nil)
return;

CRect rcClient;
GetClientRect(&rcClient);
int x = rcClient.Width() / 2;
int y = 30;

int level = get_tree_height(&rb_tree1, rb_tree1.root);

DrawRBTree(pDC, &rb_tree1, rb_tree1.root, level, x, y);
}

#define NodeDistance 24
#define Radius 20

void CTestRBTreeView::DrawRBTree(CDC *pDC, rb_tree* t, rb_tree_node* node, int level, int x, int y)
{
if(node == t->nil)
return;

CPen pen;
CPen penBlue;
penBlue.CreatePen(PS_SOLID, 1, COLORREF(0xff0000));

if(node->left != t->nil)
{
int xLeft = x - (NodeDistance * (pow(2, level) - 1)) / 4;

pDC->SelectObject(&penBlue);
pDC->MoveTo(x, y);
pDC->LineTo(xLeft, y + 60);

DrawRBTree(pDC, t, node->left, level - 1, xLeft, y + 60);

}

if(node->right != t->nil)
{
int xRight = x + (NodeDistance * (pow(2, level) - 1)) / 4;

pDC->SelectObject(&penBlue);
pDC->MoveTo(x, y);
pDC->LineTo(xRight, y + 60);

DrawRBTree(pDC, t, node->right, level - 1, xRight, y + 60);
}

if(node->color == RED)
{
pDC->SetTextColor(0x0000ff);
pen.CreatePen(PS_SOLID, 1, 0x0000ff);
}
else
{
pDC->SetTextColor(0);
pen.CreatePen(PS_SOLID, 1, COLORREF(0));
}
pDC->SelectObject(&pen);

CRect rect;
rect.left = x - 12;
rect.top = y - 12;
rect.right = x + 12;
rect.bottom = y + 12;

pDC->Ellipse(&rect);

CString str;
str.Format("%d", node->key);
if(node->key < 10)
{
pDC->TextOut(x - 4, y - 8, str);
}
else
{
pDC->TextOut(x - 8, y - 8, str);
}
}