rbtree 红黑树的 C 语言实现

rbtree 红黑树的 C 语言实现，主要是照搬了《算法导论》，红黑树是一种比较有实用意义的平衡树，不过也挻复杂，主要是利用其它兄弟节点再加以旋转来实现平衡，不过在并发上增删操作不太方便，这时用 skip list 跳跃表有一定的优势。不过用来实现key-->value 数据结构 map , dictation, set 还是不错，现在把源码和测试表如下：

/*
 * rbtree.h
 *
 *  Created on: 2018年1月16日
 *      Author: shenwei
 */

#ifndef SRC_RBTREE_H_
#define SRC_RBTREE_H_


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

struct rbtree_node{
	struct rbtree_node *left;
	struct rbtree_node *right;
	struct rbtree_node *p;
	int color;
	void *data;
};

struct rbtree{
	struct rbtree_node *root;
	struct rbtree_node *nil;
	int (*cmp)(void *a, void *b);
	void (*free)(void *data);
	int (*dump)(void *data);
};

/*
 * rbtree node api
 */

struct rbtree_node *rbtree_node_init(void *data);

int rbtree_node_exit(struct rbtree_node *node);

struct rbtree_node *rbtree_node_nil();


/*
 * rbtree api
 */

struct rbtree *rbtree_init(int cmp(void *a, void *b), void free(void *data));

int rbtree_exit(struct rbtree *root);

int rbtree_insert(struct rbtree *root, void *data);

int rbtree_delete(struct rbtree *root, void *data);

struct rbtree_node *rbtree_search_node(struct rbtree *root, struct rbtree_node *node, void *data);
struct rbtree_node *rbtree_search(struct rbtree *root, void *data);

int rbtree_dump_node(struct rbtree *root,struct rbtree_node *node, int layer);
int rbtree_dump(struct rbtree *root);

struct rbtree_node *rbtree_min(struct rbtree *root);

#endif /* SRC_RBTREE_H_ */

/*
 * rbtree.c
 *
 *  Created on: 2018年1月16日
 *      Author: shenwei
 */


#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "rbtree.h"

struct rbtree_node *rbtree_node_init(void *data){
	struct rbtree_node *node = (struct rbtree_node *)malloc(sizeof(*node));
	if(node == NULL){
		printf("malloc failed\n");
		return NULL;
	}
	memset(node, 0, sizeof(*node));
	node->data = data;
	return node;
}


struct rbtree_node *rbtree_node_nil(){
	struct rbtree_node *node = rbtree_node_init(NULL);
	if(node == NULL){
		return NULL;
	}
	node->color = BLACK;
	node->left = node;
	node->right = node;
	return node;
}

int rbtree_node_exit(struct rbtree_node *node){
	if(node){
		free(node);
	}
	return 0;
}


int rbtree_default_cmp(void *a, void *b){
	return strcmp(a, b);
}


struct rbtree *rbtree_init(int cmp(void *a, void *b), void rbfree(void *data)){
	struct rbtree *root = (struct rbtree *)malloc(sizeof(*root));
	if(root == NULL){
		printf("malloc failed\n");
		return NULL;
	}
	memset(root, 0, sizeof(*root));
	if(cmp){
		root->cmp = cmp;
	}else{
		root->cmp = rbtree_default_cmp;
	}
	if(rbfree){
		root->free = rbfree;
	}else{
		root->free = free;
	}
	root->nil = rbtree_node_nil();
	root->root = root->nil;
	return root;
}

int rbtree_free_data(struct rbtree *root, struct rbtree_node *node){
	void *data = node->data;
	root->free(data);
	node->data = NULL;
	return 0;
}


int rbtree_delete_node(struct rbtree *root, struct rbtree_node *node){

	if(root->root == root->nil){
		return 0;
	}
	rbtree_free_data(root, node);
	if(node == root->root){
		root->root = root->nil;
	}else{
		if(node->p->left == node){
			node->p->left = root->nil;
		}else{
			node->p->right = root->nil;
		}
	}
	rbtree_node_exit(node);
	return 0;
}


int rbtree_exit(struct rbtree *root){

	struct rbtree_node *node, *left, *right, *p;
	node = root->root;
	while(node != root->nil){
		left = node->left;
		right = node->right;
		if(left == root->nil && right == root->nil){
			p = node->p;
			rbtree_delete_node(root, node);
			node = p;
		}else if(left != root->nil){
			node = node->left;
		}else if(right != root->nil){
			node = node->right;
		}
	}

	rbtree_free_data(root, root->nil);
	free(root);
	return 0;
}

int rbtree_left_rotate(struct rbtree *root, struct rbtree_node *node){

	struct rbtree_node *x, *y;
	if(node == root->nil){
		printf("left node is nil\n");
		return -1;
	}
	x = node->right;
	y = x->left;

	node->right = y;
	if(y != root->nil){
		y->p = node;
	}

	x->p = node->p;
	node->p = x;
	x->left = node;

	if(root->root == node){
		root->root = x;
		x->p = root->nil;
	}else{
		if(x->p->left == node){
			x->p->left = x;
		}else{
			x->p->right = x;
		}
	}
	return 0;
}

int rbtree_right_rotate(struct rbtree *root, struct rbtree_node *node){
	struct rbtree_node *x, *y;
	if(node == root->nil){
		printf("left node is nil\n");
		return -1;
	}

	x = node->left;
	y = x->right;

	node->left = y;
	if(y != root->nil){
		y->p = node;
	}

	x->p = node->p;
	node->p = x;
	x->right = node;


	if(root->root == node){
		root->root = x;
		x->p = root->nil;
	}else{
		if(x->p->left == node){
			x->p->left = x;
		}else{
			x->p->right = x;
		}
	}
	return 0;
}

struct rbtree_node *rbtree_get_prev_node(struct rbtree *root, void *data){
	struct rbtree_node *node, *prev;
	node = root->root;
	prev = node;
	while(node != root->nil){
		prev = node;
		if(root->cmp(node->data, data) < 0){
			node = node->right;
		}else{
			node = node->left;
		}
	}
	return prev;
}

struct rbtree_node *rbtree_get_new_node(struct rbtree *root, void *data){
	struct rbtree_node *node;
	node = rbtree_node_init(data);
	node->left = root->nil;
	node->right = root->nil;
	return node;
}

int rbtree_insert(struct rbtree *root, void *data){

	struct rbtree_node *node, *new_node;
	new_node = rbtree_get_new_node(root, data);
	node = rbtree_get_prev_node(root, data);
	if(node == root->nil){
		root->root = new_node;
		root->root->color = BLACK;
		root->root->p = root->nil;
		return 0;
	}
	if(root->cmp(node->data, data) < 0){
		node->right = new_node;
	}else{
		node->left = new_node;
	}
	new_node->p = node;
	new_node->color = RED;
	if(new_node->p->color == BLACK){
		return 0;
	}

	struct rbtree_node *w;
	int old_color;
	node = new_node;
	while(node->p->color == RED){
		if(node->p->p->left == node->p){
			if(node->p->right == node){
				rbtree_left_rotate(root, node->p);
				node = node->left;
			}
			w = node->p->p->right;
			if(w->color == RED){
				w->color = BLACK;
				node->p->color = BLACK;
				node = node->p->p;
				node->color = RED;
			}else{
				w = node->p->p;
				old_color = w->color;
				rbtree_right_rotate(root, w);
				w->color = RED;
				node->p->color = old_color;
			}
		}else{
			if(node->p->left == node){
				rbtree_right_rotate(root, node->p);
				node = node->right;
			}
			w = node->p->p->left;
			if(w->color == RED){
				w->color = BLACK;
				node->p->color = BLACK;
				node = node->p->p;
				node->color = RED;
			}else{
				w = node->p->p;
				old_color = w->color;
				rbtree_left_rotate(root, w);
				w->color = RED;
				node->p->color = old_color;
			}
		}
	}
	root->root->color = BLACK;
	return 0;
}

struct rbtree_node *rbtree_min_node(struct rbtree *root, struct rbtree_node *node){

	if(node == root->nil){
		return node;
	}
	while(node->left != root->nil){
		node = node->left;
	}
	return node;
}

struct rbtree_node *rbtree_max_node(struct rbtree *root, struct rbtree_node *node){

	if(node == root->nil){
		return node;
	}
	while(node->right != root->nil){
		node = node->right;
	}
	return node;
}


int rbtree_delete_fixup(struct rbtree *root, struct rbtree_node *node){
	struct rbtree_node *w;
	while(node != root->root && node->color == BLACK){
		if(node->p->left == node){
			w = node->p->right;
			if(w->color == RED){
				w->color = BLACK;
				node->p->color = RED;
				rbtree_left_rotate(root, node->p);
				w = node->p->right; // this == prev w->left the color is BLACK
			}
			if(w->left->color == BLACK && w->right->color == BLACK){
				w->color = RED;
				node = node->p;
			}else{
				if(w->right->color == BLACK){
					w->left->color = BLACK;
					w->color = RED;
					rbtree_right_rotate(root, w);
					w = node->p->right;
				}
				w->color = node->p->color;
				node->p->color = BLACK;
				w->right->color = BLACK;
				rbtree_left_rotate(root, node->p);
				node = root->root;
			}
		}else{
			w = node->p->left;
			if(w->color == RED){
				w->color = BLACK;
				node->p->color = RED;
				rbtree_right_rotate(root, node->p);
				w = node->p->left;
			}
			if(w->left->color == BLACK && w->right->color == BLACK){
				w->color = RED;
				node = node->p;
			}else{
				if(w->left->color == BLACK){
					w->right->color = BLACK;
					w->color = RED;
					rbtree_left_rotate(root, w);
					w = node->p->left;
				}
				w->color = node->p->color;
				node->p->color = BLACK;
				w->left->color = BLACK;
				rbtree_right_rotate(root, node->p);
				node = root->root;
			}
		}
	}
	node->color = BLACK;
	return 0;
}

void rbtree_node_free(struct rbtree *root, struct rbtree_node *node){
	root->free(node->data);
	rbtree_node_exit(node);
	return ;
}

int rbtree_delete(struct rbtree *root, void *data){

	struct rbtree_node *node = rbtree_search(root, data);
	if(node == root->nil){
		printf("rbtree_delete no find the key: \n");
		if(root->dump){
			root->dump(data);
		}
		return -1;
	}

	// when node left and right child are all nil ,
	// will used nil node , this is no very good
	// in rotate option the nil node must not to modify parent point
	struct rbtree_node  *x, *y;
	if(node->left == root->nil){
		y = node->right;
		x = node;
	}else if(node->right == root->nil){
		y = node->left;
		x = node;
	}else{
		x = rbtree_min_node(root, node->right);
		if(x->left != root->nil){
			y = x->left;
		}else{
			y = x->right;
		}
	}
	if(x == root->root){
		root->root = y;
		rbtree_node_free(root, x);
		return 0;
	}

	int old_color;
	old_color = x->color;

	if(x == x->p->left){
		x->p->left = y;
	}else{
		x->p->right = y;
	}

	if(x == node){
		y->p = x->p;
	}else{
		if(x->p == node){
			y->p = x;
		}else{
			y->p = x->p;
		}


		x->left = node->left;
		x->right = node->right;
		x->p  = node->p;

		int color = x->color;
		x->color = node->color;
		node->color = color;

		if(node == root->root){
			root->root = x;
		}else{
			if(node == node->p->left){
				node->p->left = x;
			}else{
				node->p->right = x;
			}
		}
		if(x->left != root->nil){
			x->left->p = x;
		}

		if(x->right != root->nil){
			x->right->p = x;
		}
	}

	rbtree_node_free(root, node);
	if(old_color == RED){
		return 0;
	}

	rbtree_delete_fixup(root, y);
	return 0;
}


struct rbtree_node *rbtree_search_node(struct rbtree *root, struct rbtree_node *node, void *data){

	struct rbtree_node *tmp;
	if(node == root->nil){
		return root->nil;
	}

	if(root->cmp(node->data, data) == 0){
		return node;
	}

	tmp = rbtree_search_node(root, node->left,data);
	if( tmp != root->nil){
		return  tmp;
	}
	tmp = rbtree_search_node(root, node->right, data);
	if( tmp != root->nil){
		return  tmp;
	}
	return root->nil;
}

struct rbtree_node *rbtree_search(struct rbtree *root, void *data){

	return rbtree_search_node(root, root->root, data);
}

int rbtree_dump_node(struct rbtree *root,struct rbtree_node *node, int layer){

	if(node == root->nil){
		return 0;
	}

	rbtree_dump_node(root, node->left, layer+1);

	int i;
	for(i = 0; i < layer; i++){
		printf("\t");
	}
	printf("%d -- color:%s", layer,node->color?"\x1b[1;34m B \x1b[0m":"\x1b[1;31m R \x1b[0m");
	if(root->dump){
		root->dump(node->data);
	}

	rbtree_dump_node(root, node->right, layer+1);
	return 0;
}

int rbtree_dump(struct rbtree *root){
	rbtree_dump_node(root, root->root, 0);
	return 0;
}

struct rbtree_node *rbtree_min(struct rbtree *root){
	return rbtree_min_node(root, root->root);
}

/*
 * rbtree_test.c
 *
 *  Created on: 2018年1月16日
 *      Author: shenwei
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include "rbtree.h"

struct test_node{
	int key;
	char *str_key;
	void *data;
};

struct test_node *test_node_init(int key, void *data){
	struct test_node *node = (struct test_node *)malloc(sizeof(*node));
	if(node == NULL){
		printf("malloc failed\n");
		return NULL;
	}
	node->key = key;
	node->data = data;
	return node;
}

struct test_node *test_node_init_string(char *key, void *data){
	struct test_node *node = (struct test_node *)malloc(sizeof(*node));
	if(node == NULL){
		printf("malloc failed\n");
		return NULL;
	}
	node->str_key = strdup(key);
	node->data = data;
	return node;
}

void test_node_exit(void *data){
	struct test_node *node = (struct test_node *)data;
	if(node){
		if(node->data){
			free(node->data);
		}
		if(node->str_key){
			free(node->str_key);
		}
		free(node);
	}
}

int test_node_cmp(void *aa, void *bb){
	struct test_node *a = (struct test_node *)aa;
	struct test_node *b = (struct test_node *)bb;
	return a->key - b->key;
}

int test_node_cmp_string(void *aa, void *bb){
	struct test_node *a = (struct test_node *)aa;
	struct test_node *b = (struct test_node *)bb;
	return strcmp(a->str_key , b->str_key);
}

int test_node_dump(void *data){
	struct test_node *node = (struct test_node *)data;
	if(node == NULL){
		printf("input data is null\n");
		return 0;
	}
	printf("key:\x1b[0;32m%d\x1b[0m\n", node->key);
	return 0;
}

int test_node_dump_string(void *data){
	struct test_node *node = (struct test_node *)data;
	if(node == NULL){
		printf("input data is null\n");
		return 0;
	}
	printf("key:\x1b[0;32m%s\x1b[0m\n", node->str_key);
	return 0;
}

int test_int(){
	struct rbtree *root = rbtree_init(test_node_cmp, test_node_exit);
	root->dump = test_node_dump;

	int a[] = {2, 3, 9, 23, 12, 44, 6, 31, 20, 5, 19,10};
	int i;
	for(i = 0; i < sizeof(a)/sizeof(a[0]); i++){
		printf("-- add key:%d \n", a[i]);
		struct test_node *node = test_node_init(a[i], NULL);
		rbtree_insert(root, node);
		rbtree_dump(root);
		printf("---------------------------------------------\n");
	}
	rbtree_dump(root);

	int b[] = { 6, 31, 20, 5, 19,10,2, 3, 9, 23, 12, 44};
	for(i = 0; i < sizeof(b)/sizeof(b[0]); i++){
		printf("-- delete key:%d \n", b[i]);
		struct test_node *node = test_node_init(b[i], NULL);
		rbtree_delete(root, node);
		test_node_exit(node);
		rbtree_dump(root);
		printf("---------------------------------------------\n");

	}
	rbtree_exit(root);
	return 0;
}

int test_string(){

	struct rbtree *root = rbtree_init(test_node_cmp_string, test_node_exit);
		root->dump = test_node_dump_string;

		char *a[] = {"Temperatures","in","parts", "of", "the", "Siberian", "region","Yakutia",
				"have", "dropped", "to", "minus", "67", "degrees",
				"Celsius", "prompting", "even", "eyelashes", "freeze"};
		int i;
		for(i = 0; i < sizeof(a)/sizeof(a[0]); i++){
			printf("-- add key:%s \n", a[i]);
			struct test_node *node = test_node_init_string(a[i], NULL);
			rbtree_insert(root, node);
			rbtree_dump(root);
			printf("---------------------------------------------\n");
		}
		rbtree_dump(root);

		char *b[] = {"Temperatures","in","parts", "of", "the", "Siberian", "region","Yakutia",
					"have", "dropped", "to", "minus", "67", "degrees",
						"Celsius", "prompting", "even", "eyelashes", "freeze"};
		for(i = 0; i < sizeof(b)/sizeof(b[0]); i++){
			printf("-- delete key:%s \n", b[i]);
			struct test_node *node = test_node_init_string(b[i], NULL);
			rbtree_delete(root, node);
			test_node_exit(node);
			rbtree_dump(root);
			printf("---------------------------------------------\n");

		}
		rbtree_exit(root);
	return 0;
}

int main(int argc, char **argv){

	test_int();
	test_string();
	return 0;
}