红黑树的实现源码

最近因为要给ccache加入红黑树的支持, 找出来曾经实现的代码作为参考, 这才发现原来的实现都是有问题的,也怪我的测试用例写的不好, 仅仅对插入操作进行了测试, 我向所有因为阅读了这份代码而造成困惑的朋友表示道歉.

　　这次重新实现, 所有的代码推倒重新编写, 参考了linux内核中红黑树的实现算法, 并且对测试用例进行了加强,希望这是最后一个对红黑树算法的修订版本.

/*-----------------------------------------------------------
　　　 RB-Tree的插入和删除操作的实现算法
　　　 参考资料:
　　　 1) <<Introduction to algorithm>>
　　　 2) http://lxr.linux.no/linux/lib/rbtree.c

　　　 作者：http://www.cppblog.com/converse/
　　　 您可以自由的传播，修改这份代码，转载处请注明原作者

　　　 红黑树的几个性质:
　　　 1) 每个结点只有红和黑两种颜色
　　　 2) 根结点是黑色的
　　　 3)空节点是黑色的（红黑树中，根节点的parent以及所有叶节点lchild、rchild都不指向NULL，而是指向一个定义好的空节点）。
　　　 4) 如果一个结点是红色的,那么它的左右两个子结点的颜色是黑色的
　　　 5) 对于每个结点而言,从这个结点到叶子结点的任何路径上的黑色结点
　　　 的数目相同
-------------------------------------------------------------*/
　
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

typedef int key_t;
typedef int data_t;

typedef enum color_t
{
　　　 RED = 0,
　　　 BLACK = 1
}color_t;

typedef struct rb_node_t
{
　　　 struct rb_node_t *left, *right, *parent;
　　　 key_t key;
　　　 data_t data;
　　　 color_t color;
}rb_node_t;

/* forward declaration */
rb_node_t* rb_insert(key_t key, data_t data, rb_node_t* root);
rb_node_t* rb_search(key_t key, rb_node_t* root);
rb_node_t* rb_erase(key_t key, rb_node_t* root);

int main()
{
　　　 int i, count = 900000;
　　　 key_t key;
　　　 rb_node_t* root = NULL, *node = NULL;
　　　
　　　 srand(time(NULL));
　　　 for (i = 1; i < count; ++i)
　　　 {
　　　　　　　 key = rand() % count;
　　　　　　　 if ((root = rb_insert(key, i, root)))
　　　　　　　 {
　　　　　　　　　　　 printf("[i = %d] insert key %d success!n", i, key);
　　　　　　　 }
　　　　　　　 else
　　　　　　　 {
　　　　　　　　　　　 printf("[i = %d] insert key %d error!n", i, key);
　　　　　　　　　　　 exit(-1);
　　　　　　　 }

　　　　　　　 if ((node = rb_search(key, root)))
　　　　　　　 {
　　　　　　　　　　　 printf("[i = %d] search key %d success!n", i, key);
　　　　　　　 }
　　　　　　　 else
　　　　　　　 {
　　　　　　　　　　　 printf("[i = %d] search key %d error!n", i, key);
　　　　　　　　　　　 exit(-1);
　　　　　　　 }
　　　　　　　 if (!(i % 10))
　　　　　　　 {
　　　　　　　　　　　 if ((root = rb_erase(key, root)))
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 printf("[i = %d] erase key %d successn", i, key);
　　　　　　　　　　　 }
　　　　　　　　　　　 else
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 printf("[i = %d] erase key %d errorn", i, key);
　　　　　　　　　　　 }
　　　　　　　 }
　　　 }

　　　 return 0;
}

static rb_node_t* rb_new_node(key_t key, data_t data)
{
　　　 rb_node_t *node = (rb_node_t*)malloc(sizeof(struct rb_node_t));

　　　 if (!node)
　　　 {
　　　　　　　 printf("malloc error!n");
　　　　　　　 exit(-1);
　　　 }
　　　 node->key = key, node->data = data;

　　　 return node;
}

/*-----------------------------------------------------------
|　　 node　　　　　　　　　　 right
|　　 / 　　　 ==>　　　　 /
|　　 a　 right　　　　 node　 y
|　　　　　　 / 　　　　　　　　　　 /
|　　　　　　 b　 y　　　　　　　　 a　　 b
　-----------------------------------------------------------*/
static rb_node_t* rb_rotate_left(rb_node_t* node, rb_node_t* root)
{
　　　 rb_node_t* right = node->right;

　　　 if ((node->right = right->left))
　　　 {
　　　　　　　 right->left->parent = node;
　　　 }
　　　 right->left = node;

　　　 if ((right->parent = node->parent))
　　　 {
　　　　　　　 if (node == node->parent->right)
　　　　　　　 {
　　　　　　　　　　　 node->parent->right = right;
　　　　　　　 }
　　　　　　　 else
　　　　　　　 {
　　　　　　　　　　　 node->parent->left = right;
　　　　　　　 }
　　　 }
　　　 else
　　　 {
　　　　　　　 root = right;
　　　 }
　　　 node->parent = right;

　　　 return root;
}

/*-----------------------------------------------------------
|　　　　　　 node　　　　　　　　　　 left
|　　　　　　 / 　　　　　　　　　　　 /
|　　　 left　 y　　 ==>　　　 a　　 node
|　　 / 　　　　　　　　　　　　　　 /
|　 a　　 b　　　　　　　　　　　　 b　　 y
-----------------------------------------------------------*/
static rb_node_t* rb_rotate_right(rb_node_t* node, rb_node_t* root)
{
　　　 rb_node_t* left = node->left;

　　　 if ((node->left = left->right))
　　　 {
　　　　　　　 left->right->parent = node;
　　　 }
　　　 left->right = node;

　　　 if ((left->parent = node->parent))
　　　 {
　　　　　　　 if (node == node->parent->right)
　　　　　　　 {
　　　　　　　　　　　 node->parent->right = left;
　　　　　　　 }
　　　　　　　 else
　　　　　　　 {
　　　　　　　　　　　 node->parent->left = left;
　　　　　　　 }
　　　 }
　　　 else
　　　 {
　　　　　　　 root = left;
　　　 }
　　　 node->parent = left;

　　　 return root;
}

static rb_node_t* rb_insert_rebalance(rb_node_t *node, rb_node_t *root)
{
　　　 rb_node_t *parent, *gparent, *uncle, *tmp;

　　　 while ((parent = node->parent) && parent->color == RED)
　　　 {
　　　　　　　 gparent = parent->parent;

　　　　　　　 if (parent == gparent->left)
　　　　　　　 {
　　　　　　　　　　　 uncle = gparent->right;
　　　　　　　　　　　 if (uncle && uncle->color == RED)
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 uncle->color = BLACK;
　　　　　　　　　　　　　　　 parent->color = BLACK;
　　　　　　　　　　　　　　　 gparent->color = RED;
　　　　　　　　　　　　　　　 node = gparent;
　　　　　　　　　　　 }
　　　　　　　　　　　 else
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 if (parent->right == node)
　　　　　　　　　　　　　　　 {
　　　　　　　　　　　　　　　　　　　 root = rb_rotate_left(parent, root);
　　　　　　　　　　　　　　　　　　　 tmp = parent;
　　　　　　　　　　　　　　　　　　　 parent = node;
　　　　　　　　　　　　　　　　　　　 node = tmp;
　　　　　　　　　　　　　　　 }

　　　　　　　　　　　　　　　 parent->color = BLACK;
　　　　　　　　　　　　　　　 gparent->color = RED;
　　　　　　　　　　　　　　　 root = rb_rotate_right(gparent, root);
　　　　　　　　　　　 }
　　　　　　　 }
　　　　　　　 else
　　　　　　　 {
　　　　　　　　　　　 uncle = gparent->left;
　　　　　　　　　　　 if (uncle && uncle->color == RED)
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 uncle->color = BLACK;
　　　　　　　　　　　　　　　 parent->color = BLACK;
　　　　　　　　　　　　　　　 gparent->color = RED;
　　　　　　　　　　　　　　　 node = gparent;
　　　　　　　　　　　 }
　　　　　　　　　　　 else
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 if (parent->left == node)
　　　　　　　　　　　　　　　 {
　　　　　　　　　　　　　　　　　　　 root = rb_rotate_right(parent, root);
　　　　　　　　　　　　　　　　　　　 tmp = parent;
　　　　　　　　　　　　　　　　　　　 parent = node;
　　　　　　　　　　　　　　　　　　　 node = tmp;
　　　　　　　　　　　　　　　 }

　　　　　　　　　　　　　　　 parent->color = BLACK;
　　　　　　　　　　　　　　　 gparent->color = RED;
　　　　　　　　　　　　　　　 root = rb_rotate_left(gparent, root);
　　　　　　　　　　　 }
　　　　　　　 }
　　　 }

　　　 root->color = BLACK;

　　　 return root;
}

static rb_node_t* rb_erase_rebalance(rb_node_t *node, rb_node_t *parent, rb_node_t *root)
{
　　　 rb_node_t *other, *o_left, *o_right;

　　　 while ((!node || node->color == BLACK) && node != root)
　　　 {
　　　　　　　 if (parent->left == node)
　　　　　　　 {
　　　　　　　　　　　 other = parent->right;
　　　　　　　　　　　 if (other->color == RED)
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 other->color = BLACK;
　　　　　　　　　　　　　　　 parent->color = RED;
　　　　　　　　　　　　　　　 root = rb_rotate_left(parent, root);
　　　　　　　　　　　　　　　 other = parent->right;
　　　　　　　　　　　 }
　　　　　　　　　　　 if ((!other->left || other->left->color == BLACK) &&
　　　　　　　　　　　　　　　 (!other->right || other->right->color == BLACK))
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 other->color = RED;
　　　　　　　　　　　　　　　 node = parent;
　　　　　　　　　　　　　　　 parent = node->parent;
　　　　　　　　　　　 }
　　　　　　　　　　　 else
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 if (!other->right || other->right->color == BLACK)
　　　　　　　　　　　　　　　 {
　　　　　　　　　　　　　　　　　　　 if ((o_left = other->left))
　　　　　　　　　　　　　　　　　　　 {
　　　　　　　　　　　　　　　　　　　　　　　 o_left->color = BLACK;
　　　　　　　　　　　　　　　　　　　 }
　　　　　　　　　　　　　　　　　　　 other->color = RED;
　　　　　　　　　　　　　　　　　　　 root = rb_rotate_right(other, root);
　　　　　　　　　　　　　　　　　　　 other = parent->right;
　　　　　　　　　　　　　　　 }
　　　　　　　　　　　　　　　 other->color = parent->color;
　　　　　　　　　　　　　　　 parent->color = BLACK;
　　　　　　　　　　　　　　　 if (other->right)
　　　　　　　　　　　　　　　 {
　　　　　　　　　　　　　　　　　　　 other->right->color = BLACK;
　　　　　　　　　　　　　　　 }
　　　　　　　　　　　　　　　 root = rb_rotate_left(parent, root);
　　　　　　　　　　　　　　　 node = root;
　　　　　　　　　　　　　　　 break;
　　　　　　　　　　　 }
　　　　　　　 }
　　　　　　　 else
　　　　　　　 {
　　　　　　　　　　　 other = parent->left;
　　　　　　　　　　　 if (other->color == RED)
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 other->color = BLACK;
　　　　　　　　　　　　　　　 parent->color = RED;
　　　　　　　　　　　　　　　 root = rb_rotate_right(parent, root);
　　　　　　　　　　　　　　　 other = parent->left;
　　　　　　　　　　　 }
　　　　　　　　　　　 if ((!other->left || other->left->color == BLACK) &&
　　　　　　　　　　　　　　　 (!other->right || other->right->color == BLACK))
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 other->color = RED;
　　　　　　　　　　　　　　　 node = parent;
　　　　　　　　　　　　　　　 parent = node->parent;
　　　　　　　　　　　 }
　　　　　　　　　　　 else
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 if (!other->left || other->left->color == BLACK)
　　　　　　　　　　　　　　　 {
　　　　　　　　　　　　　　　　　　　 if ((o_right = other->right))
　　　　　　　　　　　　　　　　　　　 {
　　　　　　　　　　　　　　　　　　　　　　　 o_right->color = BLACK;
　　　　　　　　　　　　　　　　　　　 }
　　　　　　　　　　　　　　　　　　　 other->color = RED;
　　　　　　　　　　　　　　　　　　　 root = rb_rotate_left(other, root);
　　　　　　　　　　　　　　　　　　　 other = parent->left;
　　　　　　　　　　　　　　　 }
　　　　　　　　　　　　　　　 other->color = parent->color;
　　　　　　　　　　　　　　　 parent->color = BLACK;
　　　　　　　　　　　　　　　 if (other->left)
　　　　　　　　　　　　　　　 {
　　　　　　　　　　　　　　　　　　　 other->left->color = BLACK;
　　　　　　　　　　　　　　　 }
　　　　　　　　　　　　　　　 root = rb_rotate_right(parent, root);
　　　　　　　　　　　　　　　 node = root;
　　　　　　　　　　　　　　　 break;
　　　　　　　　　　　 }
　　　　　　　 }
　　　 }

　　　 if (node)
　　　 {
　　　　　　　 node->color = BLACK;
　　　 }

　　　 return root;
}

static rb_node_t* rb_search_auxiliary(key_t key, rb_node_t* root, rb_node_t** save)
{
　　　 rb_node_t *node = root, *parent = NULL;
　　　 int ret;

　　　 while (node)
　　　 {
　　　　　　　 parent = node;
　　　　　　　 ret = node->key - key;
　　　　　　　 if (0 < ret)
　　　　　　　 {
　　　　　　　　　　　 node = node->left;
　　　　　　　 }
　　　　　　　 else if (0 > ret)
　　　　　　　 {
　　　　　　　　　　　 node = node->right;
　　　　　　　 }
　　　　　　　 else
　　　　　　　 {
　　　　　　　　　　　 return node;
　　　　　　　 }
　　　 }

　　　 if (save)
　　　 {
　　　　　　　 *save = parent;
　　　 }

　　　 return NULL;
}

rb_node_t* rb_insert(key_t key, data_t data, rb_node_t* root)
{
　　　 rb_node_t *parent = NULL, *node;

　　　 parent = NULL;
　　　 if ((node = rb_search_auxiliary(key, root, &parent)))
　　　 {
　　　　　　　 return root;
　　　 }

　　　 node = rb_new_node(key, data);
　　　 node->parent = parent;
　　　 node->left = node->right = NULL;
　　　 node->color = RED;

　　　 if (parent)
　　　 {
　　　　　　　 if (parent->key > key)
　　　　　　　 {
　　　　　　　　　　　 parent->left = node;
　　　　　　　 }
　　　　　　　 else
　　　　　　　 {
　　　　　　　　　　　 parent->right = node;
　　　　　　　 }
　　　 }
　　　 else
　　　 {
　　　　　　　 root = node;
　　　 }

　　　 return rb_insert_rebalance(node, root);
}

rb_node_t* rb_search(key_t key, rb_node_t* root)
{
　　　 return rb_search_auxiliary(key, root, NULL);
}

rb_node_t* rb_erase(key_t key, rb_node_t *root)
{
　　　 rb_node_t *child, *parent, *old, *left, *node;
　　　 color_t color;

　　　 if (!(node = rb_search_auxiliary(key, root, NULL)))
　　　 {
　　　　　　　 printf("key %d is not exist!n");
　　　　　　　 return root;
　　　 }

　　　 old = node;

　　　 if (node->left && node->right)
　　　 {
　　　　　　　 node = node->right;
　　　　　　　 while ((left = node->left) != NULL)
　　　　　　　 {
　　　　　　　　　　　 node = left;
　　　　　　　 }
　　　　　　　 child = node->right;
　　　　　　　 parent = node->parent;
　　　　　　　 color = node->color;

　　　　　　　 if (child)
　　　　　　　 {
　　　　　　　　　　　 child->parent = parent;
　　　　　　　 }
　　　　　　　 if (parent)
　　　　　　　 {
　　　　　　　　　　　 if (parent->left == node)
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 parent->left = child;
　　　　　　　　　　　 }
　　　　　　　　　　　 else
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 parent->right = child;
　　　　　　　　　　　 }
　　　　　　　 }
　　　　　　　 else
　　　　　　　 {
　　　　　　　　　　　 root = child;
　　　　　　　 }

　　　　　　　 if (node->parent == old)
　　　　　　　 {
　　　　　　　　　　　 parent = node;
　　　　　　　 }

　　　　　　　 node->parent = old->parent;
　　　　　　　 node->color = old->color;
　　　　　　　 node->right = old->right;
　　　　　　　 node->left = old->left;

　　　　　　　 if (old->parent)
　　　　　　　 {
　　　　　　　　　　　 if (old->parent->left == old)
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 old->parent->left = node;
　　　　　　　　　　　 }
　　　　　　　　　　　 else
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 old->parent->right = node;
　　　　　　　　　　　 }
　　　　　　　 }
　　　　　　　 else
　　　　　　　 {
　　　　　　　　　　　 root = node;
　　　　　　　 }

　　　　　　　 old->left->parent = node;
　　　　　　　 if (old->right)
　　　　　　　 {
　　　　　　　　　　　 old->right->parent = node;
　　　　　　　 }
　　　 }
　　　 else
　　　 {
　　　　　　　 if (!node->left)
　　　　　　　 {
　　　　　　　　　　　 child = node->right;
　　　　　　　 }
　　　　　　　 else if (!node->right)
　　　　　　　 {
　　　　　　　　　　　 child = node->left;
　　　　　　　 }
　　　　　　　 parent = node->parent;
　　　　　　　 color = node->color;

　　　　　　　 if (child)
　　　　　　　 {
　　　　　　　　　　　 child->parent = parent;
　　　　　　　 }
　　　　　　　 if (parent)
　　　　　　　 {
　　　　　　　　　　　 if (parent->left == node)
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 parent->left = child;
　　　　　　　　　　　 }
　　　　　　　　　　　 else
　　　　　　　　　　　 {
　　　　　　　　　　　　　　　 parent->right = child;
　　　　　　　　　　　 }
　　　　　　　 }
　　　　　　　 else
　　　　　　　 {
　　　　　　　　　　　 root = child;
　　　　　　　 }
　　　 }

　　　 free(old);

　　　 if (color == BLACK)
　　　 {
　　　　　　　 root = rb_erase_rebalance(child, parent, root);
　　　 }

　　　 return root;
}