实现算法导论第三版中红黑树插入算法

#include <stdlib.h>
#include <stdio.h>
#include <assert.h>
#include <queue>
#include <iostream>

enum RBNODE_COLOR
{
	RBNODE_COLOR_RED = 0,
	RBNODE_COLOR_BLACK
};

struct RBNode
{
	RBNode * p;
	RBNode * left;
	RBNode * right;
	int      key;
	enum RBNODE_COLOR color;
	RBNode (int key);
};

RBNode::RBNode(int key)
	:p(0), left(0), right(0),
	 key(key),
	 color(RBNODE_COLOR_RED)
{}

class RBTree
{
private:
	RBNode * root;
	RBNode * nilNode;
public:
	RBTree();
	void Print();
	void Insert(RBNode * z);
private:
	void LevelOrder();
	void LeftRotate(RBNode * x);
	void RightRotate(RBNode * y);
	void InsertFixUp(RBNode * z);
};

RBTree::RBTree()
{
	nilNode = new RBNode(-1);
	nilNode->color = RBNODE_COLOR_BLACK;
	root = nilNode;
}

void RBTree::Print()
{
	LevelOrder();
	std::cout << std::endl;
}

void RBTree::Insert(RBNode * z)  
{  
	assert(z);  
	RBNode * y = nilNode;  
	RBNode * x = root;  
      
	while(x != nilNode)  
	{  
		y = x;  
          
        	if(z->key < x->key)  
           		 x = x->left;  
        	else  
            		x = x->right;  
	}  
  
	z->p = y;  
	if(y == nilNode)  
		root = z;
	else if(z->key < y->key)  
		y->left = z;  
	else  
		y->right = z;

	z->left = nilNode;
	z->right = nilNode;
	z->color = RBNODE_COLOR_RED;
	
	InsertFixUp(z);  
}

void RBTree::LevelOrder()
{
	if(root == nilNode) //root is empty
		return ;

	std::queue<RBNode *> que;
	RBNode * t = 0;
	
	que.push(root);
	
	while(!que.empty())
	{
		t = que.front();
		que.pop();
		
		if(nilNode == t->p)
			std::cout<<"Root("<<t->key<<")";
		else if(t == t->p->left)
			std::cout<<"L("<<t->key<<","<<t->p->key<<")";
		else
			std::cout<<"R("<<t->key<<","<<t->p->key<<")";
	
		if(RBNODE_COLOR_RED == t->color)
			std::cout<<"r";
		else if(RBNODE_COLOR_BLACK == t->color)
			std::cout<<"b";
		else
			std::cout<<"x";
	
		if(t->left != nilNode)
			que.push(t->left);

		if(t->right != nilNode)
			que.push(t->right);
	}
}

//assume:x.right != T.nil and root.p == T.nil.
void RBTree::LeftRotate(RBNode * x)
{
	assert(x);
	assert(x->right);
	assert(x->right != nilNode);

	RBNode * y = x->right;
	x->right = y->left; //y is x's right tree.so this assignment is safe.
	
	if(y->left != nilNode)
		y->left->p = x;

	y->p = x->p;
	
	if(x->p == nilNode) 
		root = y;
	else if(x == x->p->left)
		x->p->left = y;
	else
		x->p->right = y;

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

void RBTree::RightRotate(RBNode * y)
{
	assert(y);
	assert(y->left);
	assert(y->left != nilNode);

	RBNode * x = y->left;
	y->left = x->right;

	if(x->right != nilNode)
		x->right->p = y;

	x->p = y->p;

	if(y->p == nilNode)
		root = x;
	else if(y == y->p->left)
		y->p->left = x;
	else 
		y->p->right = x;

	x->right = y;
	y->p = x;
	
}

void RBTree::InsertFixUp(RBNode * z)
{
	//<insert first node>
	//When insert first node, the rbtree is empty.so the node is root node.
	//becaus set z's color to red in insert function, so set the color to black
	if(z == root)
	{
		z->color = RBNODE_COLOR_BLACK;
		return ;
	} 	
	//</insert first node>
	
	RBNode * y = 0;
	while(RBNODE_COLOR_RED == z->p->color)
	{
		std::cout<<"InsertFixup:while"<<std::endl;
		if(z->p == z->p->p->left) //if 1
		{
			y = z->p->p->right;
			if(y->color == RBNODE_COLOR_RED)
			{
				std::cout<<"Step:L1"<<std::endl;
				z->p->color = RBNODE_COLOR_BLACK;
				y->color = RBNODE_COLOR_BLACK;
				z->p->p->color = RBNODE_COLOR_RED;
				z = z->p->p;
			}
			else
			{
				if(z == z->p->right)
				{
					std::cout<<"Step:L2"<<std::endl;
					z = z->p;
					LeftRotate(z);
				}

				std::cout<<"Step:L3"<<std::endl;
				z->p->color = RBNODE_COLOR_BLACK;
				z->p->p->color = RBNODE_COLOR_RED;
				RightRotate(z->p->p);
			}		
		}
		else  //z->p == z->p->p->right
		{
			std::cout<<"while ... else"<<std::endl;
			
			y = z->p->p->left;
			if(y->color == RBNODE_COLOR_RED)
			{
				std::cout<<"Step::R1"<<std::endl;
				z->p->color = RBNODE_COLOR_BLACK;
				y->color = RBNODE_COLOR_BLACK;
				z->p->p->color = RBNODE_COLOR_RED;
				z = z->p->p;
			}
			else
			{
				if(z == z->p->left)
				{
					std::cout<<"Step::R2"<<std::endl;
					z = z->p;
					RightRotate(z);
				}

				std::cout<<"Step::R3"<<std::endl;
				z->p->color = RBNODE_COLOR_BLACK;
				z->p->p->color = RBNODE_COLOR_RED;
				LeftRotate(z->p->p);
			}
		} //end if 1
	} //end while

	root->color = RBNODE_COLOR_BLACK;
}

int main()
{
	RBTree tree;

	//RBNode nd11(11);
	//RBNode nd2(2);
	//RBNode n14(14);
	//tree.Insert(&nd11);
	//tree.Insert(&nd2);
	//tree.Insert(&n14);
	//tree.Print();


	//case 3
	//RBNode nd11(11);
	//RBNode nd2(2);
	//RBNode nd1(1);
	//tree.Insert(&nd11);
	//tree.Insert(&nd2);
	//tree.Insert(&nd1);
	//tree.Print();

	int arr[] = {13, 45, 32, 78, 43, 22, 90, 35, 26, 77, 1, 3, 8, 66, 39, 83, 52, 16};
	int size = sizeof(arr)/sizeof(int);
	RBNode * nd = 0;
	for(int i=0; i<size; i++)
	{
		nd = new RBNode(arr[i]);
		tree.Insert(nd);
	}

	tree.Print();
	
	return 0;
}

输出结果为：

InsertFixup:while
while ... else
Step::R2
Step::R3
InsertFixup:while
while ... else
Step::R1
InsertFixup:while
while ... else
Step::R1
InsertFixup:while
while ... else
Step::R3
InsertFixup:while
Step:L1
InsertFixup:while
Step:L2
Step:L3
InsertFixup:while
while ... else
Step::R1
InsertFixup:while
Step:L1
InsertFixup:while
Step:L1
InsertFixup:while
Step:L2
Step:L3
InsertFixup:while
Step:L3
Root(32)bL(22,32)bR(45,32)bL(3,22)rR(26,22)bL(39,45)bR(78,45)rL(1,3)bR(13,3)bL(35,39)rR(43,39)rL(66,78)bR(90,78)bL(8,13)rR(16,13)rL(52,66)rR(77,66)rL(83,90)r

由输出结果可知，测试用例已覆盖InsertFixup函数所有分支。

所构成的红黑树如下图：

此树符合红黑树的5条性质。是一颗合法的红黑树。