《算法導論》學習筆記（4）——紅黑樹（c語言實現）

#include <stdlib.h>
#include <stdio.h>

typedef int EleType;

typedef enum Color  //顏色屬性：紅、黑
{
	RED = 0, BLACK = 1
} Color;

typedef struct Node
{
	struct Node* left;
	struct Node* right;
	struct Node* parent;
	Color color;
	EleType key;
} Node;

//哨兵結點，表示空，黑色。紅黑樹中的所有指向空的指針均指向該結點
Node tree_nil = { &tree_nil, &tree_nil, &tree_nil, BLACK, 0 };

void rb_display( Node* root );
Node* rb_search( Node* root, int key );
Node* tree_maximum( Node* x );
Node* tree_minimum( Node* x );
Node* tree_successor( Node* x );
Node* left_rotate( Node* root, Node* x );
Node* right_rotate( Node* root, Node* x );
Node* rb_insert( Node* root, int key );
Node* rb_insertFixup(Node* root, Node* x );
Node* rb_delete( Node* root, EleType key );
Node* rb_deleteFixup( Node* root, Node* x );

//中序遍歷紅黑樹，輸出所有結點及其顏色。遞歸實現
void rb_display( Node* root )
{
	if( root != &tree_nil )
	{
		rb_display( root->left );
		printf( "%d, color is %s\n", root->key, (root->color?"RED":"BLACK") );
		rb_display( root->right );
	}
 }

//查找結點。給定關鍵字，返回其所在結點。若紅黑樹中不存在該結點，返回空。遞歸實現
Node* rb_search( Node* root, int key )
{
	if( (root == &tree_nil) || (root->key == key) )
		return root;
	if( key < root->key )
		return rb_search( root->left, key );
	else
		return rb_search( root->right, key );
 }

//返回結點的最大子結點，遞歸實現
Node* tree_maximum( Node* x )
{
	if( x->right == &tree_nil )
		return x;
	else
		return tree_maximum( x->right );
}

//返回結點的最小子結點，遞歸實現
Node* tree_minimum( Node* x )
{
	if( x->left == &tree_nil )
		return x;
	else
		return tree_minimum( x->left );
}

//返回結點的後繼（大於x.key的最小結點）
Node* tree_successor( Node* x )
{
	if( x->right != &tree_nil )
		return tree_minimum( x->right );
	Node* y = x->parent;
	while( (y != &tree_nil) && (x == y->right) )
	{
		x = y;
		y = y->parent;
	}
	return y;
}

//左旋
Node* left_rotate( Node* root, Node* x )
{
	Node* y;
	if( x->right == &tree_nil ) //左旋結點必須有右子結點
	{
	    printf( "have no right child,rotation cancel.\n" );
	    return root;
	}
	y = x->right;
	x->right = y->left; //β子樹從y的左子樹變成x的右子樹
	if( y->left != &tree_nil )
		y->left->parent = x;
	y->parent = x->parent; //y的父結點從x變爲x的父結點
	if( x->parent == &tree_nil )
		root = y;
	else if( x == x->parent->left )
		x->parent->left = y;
	else
		x->parent->right = y;
	y->left = x;
	x->parent = y;
	return root;
}

//右旋
Node* right_rotate( Node* root, Node* x )
{
	Node* y;
	if( x->left == &tree_nil ) //右旋結點必須有左子結點
	{
	    printf( "have no left child,rotation cancel.\n" );
	    return root;
	}
	y = x->left;
	x->left = y->right;
	if( y->right != &tree_nil )
		y->right->parent = x;
	y->parent = x->parent;
	if( x->parent == &tree_nil )
		root = y;
	else if( x == x->parent->left )
		x->parent->left = y;
	else
		x->parent->right = y;
	y->right = x;
	x->parent = y;
	return root;
}

Node* rb_insert( Node* root, int key )
{
	Node* y = &tree_nil;
	Node* x = root;
	Node* z = (Node*) malloc( sizeof(Node) );
	z->key = key;
	while( x != &tree_nil )
	{
		y = x;
		if ( key < x->key )
			x = x->left;
		else
			x = x->right;
	}
	z->parent = y;
	if ( y == &tree_nil )
		root = z;
	else if ( z->key < y->key )
		y->left = z;
	else
		y->right = z;
	z->left = &tree_nil;
	z->right = &tree_nil;
	z->color = RED;
	return rb_insertFixup( root, z );
}

Node* rb_insertFixup( Node* root, Node* z )
{
	while( z->parent->color == RED )
	{
		if( z->parent == z->parent->parent->left )
		{
			Node* 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;
					root = left_rotate( root, z );
				}
				z->parent->color = BLACK; //第三種情況：當前結點是左孩子，且叔結點是黑色
				z->parent->parent->color = RED;
				root = right_rotate( root, z->parent->parent );
			}
		}
		else
		{
			Node* 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;
					root = right_rotate( root, z );
				}
				z->parent->color = BLACK;
				z->parent->parent->color = RED;
				root = left_rotate( root, z->parent->parent );
			}
		}
	}
	root->color = BLACK;
	return root;
}

Node* rb_delete( Node* root, EleType key )
{
	Node* z;
	z = rb_search( root, key );
	if( z == &tree_nil )
	{
		printf("The element %d is not in the tree\n", key );
		return root;
	}
	else
	{
		Node *x, *y;
		if( (z->left == &tree_nil) || (z->right == &tree_nil) )
			y = z;
		else
			y = tree_successor(z);
		if ( y->left != &tree_nil )
			x = y->left;
		else
			x = y->right;
		x->parent = y->parent;
		if ( y->parent == &tree_nil )
			root = x;
		else if ( y == x->parent->left )
			y->parent->left = x;
		else
			y->parent->right = x;
		if ( y != z )
			z->key = y->key;
		if ( y->color == BLACK )
			root = rb_deleteFixup( root, x );
		free(y);
		return root;
	}
}

Node* rb_deleteFixup( Node* root, Node* x )
{
	Node* w;
	while ( (x != root) && (x->color == BLACK) )
	{
		if ( x == x->parent->left )
		{
			w = x->parent->right;
			if (w->color == RED) //情況1；待刪結點的兄弟結點w是紅色
			{
				w->color = BLACK;
				x->parent->color = RED;
				root = left_rotate( root, x->parent );
				w = x->parent->right;
			}
			if ( (w->left->color == BLACK) &&
				 (w->right->color == BLACK) ) //情況2；待刪結點的兄弟結點w是黑色，且w的兩個子結點都是黑色
			{
				w->color = RED;
				x = x->parent;
			}
			else
			{
				if ( w->right->color == BLACK ) //情況3；待刪結點的兄弟結點w是黑色，且w左孩子紅色，右孩子黑色
				{
					w->left->color = BLACK;
					w->color = RED;
					root = right_rotate( root, w );
					w = x->parent->right;
				}
				w->color = x->parent->color; //情況4；待刪結點的兄弟結點w是黑色，且w的右孩子是紅色
				x->parent->color = BLACK;
				w->right->color = BLACK;
				root = left_rotate( root, x->parent );
				x = root;
			}

		}
		else
		{
			w = x->parent->left;
			if ( w->color == RED )
			{
				w->color = BLACK;
				x->parent->color = RED;
				root = right_rotate( root, x->parent );
				w = x->parent->right;
			}
			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;
					root = left_rotate( root, w );
					w = x->parent->left;
				}
				w->color = x->parent->color;
				x->parent->color = BLACK;
				w->left->color = BLACK;
				root = right_rotate( root, x->parent );
				x = root;
			}

		}
	}
	x->color = BLACK;
	return root;
}

int main()
{
	Node* root;
	root = &tree_nil;
	root->parent = &tree_nil;
	root = rb_insert( root, 11 );
	root = rb_insert( root, 1 );
	root = rb_insert( root, 9 );
	root = rb_insert( root, 2 );
	root = rb_insert( root, 10 );
	rb_display( root );

	root = rb_delete( root, 2 );
	rb_display( root );

	root = rb_delete( root, -10 );
	rb_display( root );

	return 0;
}