算法导论代码 第13章 红黑树

#include <stdio.h>
#include <string.h>
#include <stdlib.h>
typedef struct red_black_tree_type *tree;
enum color_enum {
	color_red,
	color_black
};
struct tree_node {
	void *key;
	enum color_enum color;
	struct tree_node *parent;
	struct tree_node *left;
	struct tree_node *right;
};

struct red_black_tree_type {
	int (*comp) (const void *, const void *);
	struct tree_node *root;
	struct tree_node *nil;
};
void tree_node_ini(struct tree_node *p, void *key)
{
	p->key = key;
	p->parent = NULL;
	p->left = NULL;
	p->right = NULL;
	p->color = color_black;
}

void tree_left_rotate(tree t, struct tree_node *x)
{
	struct 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 tree_right_rotate(tree t, struct tree_node *x)
{
	struct 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->left) {
			x->parent->left = y;
		} else {
			x->parent->right = y;
		}
	}
	y->right = x;
	x->parent = y;
}

tree tree_create(int (*comp) (const void *, const void *))
{
	tree t = malloc(sizeof(struct red_black_tree_type));
	t->comp = comp;
	t->nil = malloc(sizeof(struct tree_node));
	tree_node_ini(t->nil, NULL);
	t->root = t->nil;
	return t;
}

void tree_destroy_all_node(tree t, struct tree_node *x,
			   void (*free_key) (void *))
{
	if (x != t->nil) {
		tree_destroy_all_node(t, x->left, free_key);
		tree_destroy_all_node(t, x->right, free_key);
		free_key(x->key);
		free(x);
	}
}

void tree_destroy(tree t, void (*free_key) (void *))
{
	tree_destroy_all_node(t, t->root, free_key);
	free(t->nil);
	free(t);
}

void tree_inorder_walk(tree t, struct tree_node *x,
		       void (*handle) (const void *))
{
	if (x != t->nil) {
		tree_inorder_walk(t, x->left, handle);
		handle(x->key);
		printf(" ");
		tree_inorder_walk(t, x->right, handle);
	}
}

struct tree_node *tree_search(tree t, struct tree_node *x, void *key)
{
	if (x == t->nil || t->comp(key, x->key) == 0) {
		return x;
	}
	if (t->comp(key, x->key) < 0) {
		return tree_search(t, x->left, key);
	} else {
		return tree_search(t, x->right, key);
	}
}

void tree_count_leaf(tree t, struct tree_node *x,
		     struct tree_node *leaf_array[], int *num)
{
	if (x != t->nil) {
		if (x->left == t->nil && x->right == t->nil) {
			leaf_array[++*num] = x;
		}
		tree_count_leaf(t, x->left, leaf_array, num);
		tree_count_leaf(t, x->right, leaf_array, num);
	}
}

void tree_black_height(tree t, struct tree_node *x,
		       struct tree_node *leaf_node,
		       int *black_height)
{
	if (x->color == color_black) {
		++*black_height;
	}
	if (x == t->nil) {
		return;
	}
	if (t->comp(leaf_node->key , x->key)<0) {
		tree_black_height(t, x->left, leaf_node, black_height);
	} else {
		tree_black_height(t, x->right, leaf_node, black_height);
	}
}

struct tree_node *tree_minimum(tree t, struct tree_node *x)
{
	while (x != t->nil && x->left != t->nil) {
		x = x->left;
	}
	return x;
}

struct tree_node *tree_successor(tree t, struct tree_node *x)
{
	if (x->right != t->nil) {
		return tree_minimum(t, x->right);
	}
	struct tree_node *y = x->parent;
	while (y != t->nil && x == y->right) {
		x = y;
		y = y->parent;
	}
	return y;
}

void tree_insert_fixup(tree t, struct tree_node *z)
{
	struct tree_node *y = t->nil;
	while (z->parent->color == color_red) {
		if (z->parent == z->parent->parent->left) {
			y = z->parent->parent->right;
			//情况1:z的叔叔y是红色的
			if (y->color == color_red) {
				z->parent->color = color_black;
				y->color = color_black;
				z->parent->parent->color = color_red;
				z = z->parent->parent;
			} else {
				//情况2:z的叔叔y是黑色的,z是右孩子
				if (z == z->parent->right) {
					z = z->parent;
					tree_left_rotate(t, z);
				}
				//情况3:z的叔叔y是黑色的,z是左孩子
				z->parent->color = color_black;
				z->parent->parent->color = color_red;
				tree_right_rotate(t, z->parent->parent);
			}
		} else {
			y = z->parent->parent->left;
			if (y->color == color_red) {
				z->parent->color = color_black;
				y->color = color_black;
				z->parent->parent->color = color_red;
				z = z->parent->parent;
			} else {
				if (z == z->parent->left) {
					z = z->parent;
					tree_right_rotate(t, z);
				}
				z->parent->color = color_black;
				z->parent->parent->color = color_red;
				tree_left_rotate(t, z->parent->parent);
			}
		}
	}
	t->root->color = color_black;
}

void tree_insert(tree t, struct tree_node *z)
{
	struct tree_node *y = t->nil;
	struct tree_node *x = t->root;
	while (x != t->nil) {
		y = x;
		if (t->comp(z->key, x->key) < 0) {
			x = x->left;
		} else {
			x = x->right;
		}
	}
	z->parent = y;
	if (y == t->nil) {
		t->root = z;
	} else {
		if (t->comp(z->key, y->key) < 0) {
			y->left = z;
		} else {
			y->right = z;
		}
	}
	z->left = t->nil;
	z->right = t->nil;
	z->color = color_red;
	tree_insert_fixup(t, z);
}

void tree_delete_fixup(tree t, struct tree_node *x)
{
	struct tree_node *w = t->nil;
	while (x != t->root && x->color == color_black) {
		if (x == x->parent->left) {
			w = x->parent->right;
			//情况1:x的兄弟w是红色的
			if (w->color == color_red) {
				w->color = color_black;
				x->parent->color = color_red;
				tree_left_rotate(t, x->parent);
				w = x->parent->right;
			} else {
				//情况2:x的兄弟w是黑色的,而且w的两个孩子都是黑色的
				if (w->left->color == color_black
				    && w->right->color == color_black) {
					w->color = color_red;
					x = x->parent;
				} else {
					//情况3:x的兄弟w是黑色的,w的左孩子是红的，右孩子是黑的
					if (w->right->color == color_black) {
						w->left->color = color_black;
						w->color = color_red;
						tree_right_rotate(t, w);
						w = x->parent->right;
					}
					//情况4:x的兄弟w是黑色的,而且w的右孩子是红色的
					w->color = x->parent->color;
					x->parent->color = color_black;
					w->right->color = color_black;
					tree_left_rotate(t, x->parent);
					x = t->root;
				}
			}
		} else {
			w = x->parent->left;
			if (w->color == color_red) {
				w->color = color_black;
				x->parent->color = color_red;
				tree_right_rotate(t, x->parent);
				w = x->parent->left;
			} else {
				if (w->right->color == color_black
				    && w->left->color == color_black) {
					w->color = color_red;
					x = x->parent;
				} else {
					if (w->left->color == color_black) {
						w->right->color = color_black;
						w->color = color_red;
						tree_left_rotate(t, w);
						w = x->parent->left;
					}
					w->color = x->parent->color;
					x->parent->color = color_black;
					w->left->color = color_black;
					tree_right_rotate(t, x->parent);
					x = t->root;
				}
			}
		}
	}
	x->color = color_black;
}

void swap(void *a, void *b, size_t elem_size)
{
	if (a == NULL || b == NULL || a == b)
		return;
	char temp[elem_size];	/*变长数组 */
	memcpy(temp, a, elem_size);
	memcpy(a, b, elem_size);
	memcpy(b, temp, elem_size);
}

struct tree_node *tree_delete(tree t, struct tree_node *z)
{
	struct tree_node *y;
	struct tree_node *x;
	if (z->left == t->nil || z->right == t->nil) {
		y = z;
	} else {
		y = tree_successor(t, z);
	}
	if (y->left != t->nil) {
		x = y->left;
	} else {
		x = y->right;
	}
	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;
		}
	}
	if (y != z) {
		/*要删除的结点y是z的后继,交换z和y结点的内容 */
		swap(&z->key, &y->key, sizeof(void *));
	}
	if (y->color == color_black) {
		tree_delete_fixup(t, x);
	}
	return y;
}

int cmp_int(const void *p1, const void *p2)
{
	const int *pa = p1;
	const int *pb = p2;
	if (*pa < *pb)
		return -1;
	if (*pa == *pb)
		return 0;
	return 1;
}
void print_key(const void *key)
{
	const int *p = key;
	printf("%-2d", *p);
}

int main()
{
	tree t = tree_create(cmp_int);
	for (int i = 0; i < 10; i++) {
		struct tree_node *node = malloc(sizeof(struct tree_node));
		int *ip = malloc(sizeof(int));
		*ip = i;
		tree_node_ini(node, ip);
		tree_insert(t, node);
	}
	printf("中序遍历结果:\n");
	tree_inorder_walk(t, t->root, print_key);
	printf("\n");
	int k = 0;
	struct tree_node *result = tree_search(t, t->root, &k);
	printf("查找关键字:%d的结果:%s\n", k,
	       result != t->nil ? "成功" : "失败");
	struct tree_node *del_node = tree_delete(t, result);
	free(del_node->key);
	free(del_node);
	result = tree_search(t, t->root, &k);
	printf("删除关键字:%d的结果:%s\n", k,
	       result == t->nil ? "成功" : "失败");
	printf("查看黑高度:\n");
	struct tree_node *left_array[100];
	int num = -1;
	/*统计叶子结点 */
	tree_count_leaf(t, t->root, left_array, &num);
	for (int i = 0; i < num; i++) {
		/*黑高度不包括自身,所以初始值为-1 */
		int black_height = -1;
		tree_black_height(t, t->root, left_array[i], &black_height);
		/*测试黑高度，从根到叶的路径上，黑结点数是一样的 */
		printf("根到叶子:%d的黑高度:%d\n",
		       *(int *)left_array[i]->key, black_height);
	}
	printf("中序遍历结果:\n");
	tree_inorder_walk(t, t->root, print_key);
	printf("\n");
	/*遍历树，释放结点的key */
	tree_destroy(t, free);
	return 0;
}