红黑树的实现

学习内核调度，2.6.23以后用了一种新的调度算法，红黑树，听起来好像很强大的样子，

查资料，发现是二叉树的变种之一，看原理不是很明白的，就想自己写个代码，帮助理解，

代码很简陋，可优化的地方还有很多，仅供参考。

代码里面的树形输出还有些问题

// rbtree.cpp : 定义控制台应用程序的入口点。
//

#include "stdafx.h"
#include <stdio.h>
#include <stdlib.h>
#include <iostream>
#include <math.h>

using namespace  std;

typedef struct retree
{
	int color;
	int key;
	struct retree * left;
	struct retree * right;
	struct retree * parent;
}*rbtree;

typedef struct treeshow
{
	int level;
	int color;
	int key;
	int ser_num;
	int pkey;
}Tshow;

#define  RED 0
#define BLACK 1

int k = 0;
int flag = 0;
struct treeshow nod[100];
rbtree list = NULL;

//收集树的结点，并前序遍历编号，确定结点所在的深度
//list,待收集的树指针
//n,结点的层
void PrintTree(rbtree list, int n)
{ 
	if(list==NULL) 
		return;
	if (list->left)
	{
		PrintTree(list->left, n+1);
	}
	//printf("%d", list->key);
	nod[k].level = n;
	nod[k].key = list->key;
	nod[k].color = list->color;
	nod[k].ser_num = k + 1;
	if (list->parent)
	{
		nod[k].pkey = list->parent->key;
	}
	else nod[k].pkey = 0;
	k++;
	if (list->right)
	{
		PrintTree(list->right, 1+n);
	}
}

//计算树的高度
//t,待计算的树
int height(rbtree t)
{
	int lh,rh,h;
	if(t==NULL) return -1;
	lh=height(t->left);
	rh=height(t->right);
	h=lh>rh?lh:rh;
	return (h+1);
}

//修正树的颜色
//pt, 待修正的树指针
void change_color(rbtree pt)
{
	if (BLACK == pt->parent->color)
	{
		pt->color = RED;

	}
	if (RED == pt->parent->color)
	{
		pt->color = BLACK;
	}
	if (pt->left)
	{
		change_color(pt->left);
	}
	if (pt->right)
	{
		change_color(pt->right);
	}
}

//树右旋
//pt，待旋转的树
//n, 树的高度差
void right_rotate(rbtree pt, int n)
{
	rbtree ptree = pt;
	while(n > 1)
	{
		rbtree parent = ptree->parent;
		rbtree left = ptree->left;
		ptree->left->parent = ptree->parent;
		ptree->left = left->right;
		left->right = ptree;
		ptree->parent = left;
		parent->left = left;
		ptree = left;
		n -= 2;
	}
	change_color(ptree);
}

//树左旋
//pt,待旋转的数
//n,左右子树高度差
void left_rotate(rbtree pt, int n)
{
	rbtree ptree = pt;
	while(1)
	{
		rbtree parent = ptree->parent;
		rbtree ri = ptree->right;
		ri->parent = parent;
		ptree->right = ri->left;
		ri->left = ptree;
		ptree->parent = ri;
		ptree = ptree->right;
		parent->right = ri;
		ptree = ri;
		n -= 2;
	}
	change_color(ptree);
}

//平衡数
//t，要平衡的树的指针
int btree(rbtree t)
{
	int lh,rh,h;
	int hflag = 0;
	if(t==NULL) return -1;
	lh=btree(t->left);
	rh=btree(t->right);
	if ( lh >= rh + 2 && rh >= 0)
	{
		right_rotate(t, lh - rh);
		hflag++;
	}
	if (rh >= lh + 2 && lh >= 0)
	{
		left_rotate(t, rh - lh);
		hflag++;
	}
	if (hflag)
	{
		lh=btree(t->left);
		rh=btree(t->right);
	}
	h=lh>rh?lh:rh;
	return (h+1);
}

//建立红黑树
//key_group, 输入的原始数据，有序，递增
//n,原始数据数量
//color,结点颜色，为递归建立红黑树时使用，在main函数里调用是无效，因为首先建立的结点是根节点，必定为黑色
//po,建立结点的父节点，为递归时使用，首次调用时为NULL
rbtree build_tree(int * key_group, int n, int color, rbtree po)
{
	if (n == 0 || *key_group == NULL)
	{
		printf("the data input had error\n");
		return NULL;
	}
	//printf("n:%d\n", n);
	rbtree parent = NULL;
	if (n == 1)
	{
		{
			parent = (rbtree)malloc(sizeof(retree));
			parent->color = color;
			parent->key = *(key_group);
			//printf("su:%d\n",parent->key);
			parent->left = NULL;
			parent->right = NULL;
			parent->parent = po;
			return parent;
		}
	}
	if (n > 1)
	{
		int * key = key_group;
		int root = n/2;
		parent = (rbtree)malloc(sizeof(retree));
		if ( flag == 0)
		{
			parent->color = BLACK;
			flag = 1;
		}
		else parent->color = color;
		parent->key = *(key + root);
		//printf("su:%d\n",parent->key);
		parent->left = NULL;
		parent->right = NULL;
		parent->parent = po;
		if (root == 1)
		{
			if (n == 2)
			{
				if (BLACK == parent->color)
				{
					parent->left = build_tree(key_group, 1, RED, parent);
				}
				else parent->left = build_tree(key_group, 1, BLACK, parent);
			}
			if (n == 3)
			{
				if (BLACK == parent->color)
				{
					parent->left = build_tree(key_group, 1, RED, parent);
					parent->right = build_tree(key_group+2, 1, RED, parent);
				}
				else 
				{
					parent->left = build_tree(key_group, 1, BLACK, parent);
					parent->right = build_tree(key_group+2, 1, BLACK, parent);
				}
			}
		}
		else
		{
			if (BLACK == parent->color)
			{
				parent->left = build_tree(key_group, root, RED, parent);
				parent->right = build_tree(key_group+root+1, n - root-1, RED, parent);
			}
			else
			{
				parent->left = build_tree(key_group, root, BLACK, parent);
				parent->right = build_tree(key_group+root+1, n - root-1, BLACK, parent);
			}
		}
	}
	
	return parent;
}

//添加结点
//key,添加结点的值
void add_point(int key)
{
	int pos;
	rbtree pt = list;
	while(1)
	{
		if (key <= pt->key)
		{
			if (pt->left)
			{
				pt = pt->left;
			}
			else
			{
				pos = 0;
				break;
			}
		}
		if (key > pt->key)
		{
			if (pt->right)
			{
				pt = pt->right;
			}
			else
			{
				pos = 1;
				break;
			}
		}
	}
	rbtree tmp = (rbtree)malloc(sizeof(retree));
	tmp->color = RED;
	tmp	->key = key;
	tmp->left = NULL;
	tmp->right = NULL;
	tmp->parent = NULL;
	//if (NULL == pt->parent)
	{
		printf("pt.key:%d\n", pt->key);
		if (!pos)
		{
			pt->left = tmp;
			tmp->parent =pt;
		}
		else
		{
			pt->right = tmp;
			tmp	->parent = pt;
		}
		if (RED == pt->color)
		{
			pt->color = BLACK;
		}
	}
	btree(list);
}
//删除结点
//key,欲删除结点的值
void del_point(int key)
{
	rbtree pt = list;
	int dkey;
	int dflag;
	while(1)
	{
		if (key == pt->key)
		{
			break;
		}
		if (key > pt->key)
		{
			pt = pt->right;
		}
		if (key < pt->key)
		{
			pt = pt->left;
		}
	}
	if (!pt->parent)
	{
		if (!pt->left && !pt->right)
		{
			list = NULL;
		}
		if (pt->left && !pt->right)
		{
			list = pt->left;
			pt->left->parent = NULL;
		}
		if (!pt->left && pt->right)
		{
			list = pt->right;
			pt->right->parent = NULL;
		}
		if (pt->left && pt->right)
		{
			rbtree pm = pt->left;
			while(1)
			{
				if (pm->right)
				{
					pm = pm->right;
				}
				else break;
			}
			dkey = pm->key;
			pt->key = dkey;
			if (pm->parent->left->key == dkey)
			{
				pm->left->parent = pt;
				pt->left = pm->left;
			}
			if (pm->parent->right->key == dkey)
			{
				if (pm->left)
				{
					pm->left->parent =pm->parent;
					pm->parent ->right =pm->left;
				}
				else pm->parent->right = NULL;
			}
			pt = pm;
		}
		free(pt);
	}
	else
	{
		if (pt->parent->left->key == key)
		{
				dflag = 0;
		}
		if (pt->parent->right->key == key)
		{
			dflag = 1;
		}
	}
	if (NULL == pt->left && NULL == pt->right)
	{
		if (!dflag)
		{
			pt->parent->left = NULL;
		}
		else pt->parent->right = NULL;
		free(pt);
	}
	if (pt->left && NULL == pt->right)
	{
		if (!dflag)
		{
			pt->parent->left = pt->left;
			pt->left->parent = pt->parent;
		}
		else
		{
			pt->parent->right = pt->left;
			pt->left->parent = pt->parent;
		}
	}
	if (!pt->left && pt->right)
	{
		if (!dflag)
		{
			pt->parent->left = pt->right;
			pt->right->parent = pt->parent;
		}
		else
		{
			pt->parent->right = pt->right;
			pt->right->parent =pt->parent;
		}

	}
	if (pt->left && pt->right)
	{
		rbtree pm = pt->left;
		while(1)
		{
			if (pm->right)
			{
				pm = pm->right;
			}
			else break;
		}
		if (pm->left)
		{
			dkey = pm->key;
			pm->left->parent = pm->parent;
			pm->parent->right = pm->left;
		}
		else 
		{
			dkey = pm->key;
			pm->parent->right =NULL;
		}
		pt->key = dkey;
		pt = pm;
	}
	free(pt);
	btree(list);
}
//释放数
//list, 数的指针
void free_tree(rbtree list)
{
	if (list)
	{
		if (list->left)
		{
			free_tree(list->left);
			list->left = NULL;
		}
		if (list->right)
		{
			free_tree(list->right);
			list->right = NULL;
		}
		free(list);
		list = NULL;
	}
}
//将红黑树树形输出
//tree_size, 数结点的数量
//hei,数的高度
void show_tree(int tree_size, int hei)
{
	struct treeshow * pt;
	int i;
	int base = 10;
	int len = 4;
	int ltree = pow((float)2, hei);
	int blank;
	int num = 0;
	printf("--------------------------------------------------------------------------\n");
	while(num <= hei)
	{
		pt = nod;
		i = 0;
		for(i=0; i < base; i++) printf(" ");
		i = 0;
		blank = 0;
		while(i < tree_size)
		{
			if (pt->level == num)
			{
				for (int j = 1; j < (pt->ser_num - blank); j++)
				{
					printf("    ");
				}
				blank = pt->ser_num;
				printf("%d|", pt->key);
				if(pt->color)
					printf("b%d", pt->pkey);
				else printf("r%d",pt->pkey);
			}
			pt++;
			i++;
		}
		num++;
		printf("\n");
	}
	printf("--------------------------------------------------------------------------\n");
}

//红黑树输出函数
//输出list所有结点
void showpoint()
{
	k = 0;
	PrintTree(list, 0);
	int le = height(list);
	printf("hei:%d, k:%d\n", le, k);
	show_tree(k, le);
}

int _tmain(int argc, _TCHAR* argv[])
{
	int arr[] = {1,4,6,7,8,10,11,13,14};
	int arr_len = sizeof(arr)/sizeof(arr[0]);
	printf("the arr len:%d\n", arr_len);
	list = build_tree(arr, arr_len, 1, NULL);
	add_point(3);
	showpoint();

	del_point(4);
	showpoint();

	free_tree(list);
	return 0;
}