数据结构之——红黑树

#pragma once

#include <iostream>
using namespace std;

//结点的颜色 红or黑
enum Color
{
	RED,
	BLACK
};

//结点结构体
template <class K, class V>
struct RBTreeNode
{
	K _key;
	V _val;
	RBTreeNode<K, V>* _left;
	RBTreeNode<K, V>* _right;
	RBTreeNode<K, V>* _parent;
	Color _col;

	RBTreeNode(const K& key, const V& val)
		:_key(key)
		,_val(val)
		,_left(NULL)
		,_right(NULL)
		,_parent(NULL)
		,_col(RED)
	{}
};

//红黑树类
template <class K, class V>
class RBTree
{
	typedef RBTreeNode<K, V> Node;

public:
	RBTree()
		:_root(NULL)
	{}
	~RBTree()
	{
		_Clear(_root);
	}

	//插入结点
	bool Insert(const K& key, const V& val)
	{
		if(_root == NULL)//如果根结点为NULL，创建新的结点为根结点返回true
		{
			_root = new Node(key, val);
			_root->_col = BLACK;
			return true;
		}

		//如果根结点不为NULL，则遍历树直到找到合适的插入位置
		Node* cur = _root;
		Node* prev = cur;
		while(cur != NULL)
		{
			if(cur->_key == key)//如果树中已经有该结点，则返回false
				return false;
			else if(cur->_key > key)//如果关键值小于结点的关键值，则去左子树找
			{
				prev = cur;
				cur = cur->_left;
			}
			else//否则关键值大于结点的关键值，去右子树找
			{
				prev = cur;
				cur = cur->_right;
			}
		}
		//循环结束，找到了合适的插入位置，判断应该插入到结点的左边还是右边
		Node* tmp = new Node(key, val);
		if(prev->_key > key)
			prev->_left = tmp;
		else
			prev->_right = tmp;
		tmp->_parent = prev;

		//插入结点之后，就要开始判断目前的树是否符合红黑树的性质
		while((tmp != _root) && (prev->_col == RED))
		{
			Node* grandfather = prev->_parent;//提取出tmp的祖父结点
			if(grandfather->_left == prev)//如果prev是grandfather的左结点
			{
				//第一种情况
				//tmp为红，prev为红，grandfather为黑，uncle存在且为红
				//则将prev,uncle改为黑，grandfather改为红，然后把grandfather当成tmp，继续向上调整。
				Node* uncle = grandfather->_right;
				if(uncle != NULL && uncle->_col == RED)
				{
					prev->_col = BLACK;
					uncle->_col = BLACK;
					grandfather->_col = RED;

					tmp = grandfather;
					prev = tmp->_parent;
				}
				else//第二种情况:tmp为红，prev为红，grandfather为黑，uncle不存在/uncle为黑
					//prev为grandfather的左孩子，tmp为prev的左孩子，则进行右单旋转；
					//prev、grandfather变色--prev变黑，grandfather变红

				{
					//第三种情况
					//tmp为红，prev为红，grandfather为黑，uncle不存在/uncle为黑
					//prev为grandfather的左孩子，tmp为prev的右孩子，则针对prev做左单旋转；
					//则转换成了情况二

					if(prev->_right == tmp)
					{
						_RotateL(prev);
						swap(tmp, prev);
					}
					_RotateR(grandfather);//进行右单旋
					prev->_col = BLACK;
					grandfather->_col = RED;
					break;
				}
			}
			else//当perv是grandfather的右结点的时候，和上面的情况相反
			{
				//第一种情况
				//tmp为红，prev为红，grandfather为黑，uncle存在且为红
				//则将prev,uncle改为黑，grandfather改为红，然后把grandfather当成tmp，继续向上调整。
				Node* uncle = grandfather->_left;
				if(uncle != NULL && uncle->_col == RED)
				{
					prev->_col = BLACK;
					uncle->_col = BLACK;
					grandfather->_col = RED;

					tmp = grandfather;
					prev = tmp->_parent;
				}
				else//第二种情况:tmp为红，prev为红，grandfather为黑，uncle不存在/uncle为黑
					//prev为grandfather的右孩子，tmp为prev的右孩子，则进行左单旋转
					//prev、grandfather变色--prev变黑，grandfather变红

				{
					//第三种情况
					//tmp为红，prev为红，grandfather为黑，uncle不存在/uncle为黑
					//prev为grandfather的右孩子，tmp为prev的左孩子，则针对prev做右单旋转
					//则转换成了情况二

					if(prev->_left == tmp)
					{
						_RotateR(prev);
						swap(tmp, prev);
					}
					_RotateL(grandfather);//进行右单旋
					prev->_col = BLACK;
					grandfather->_col = RED;
					break;
				}
			}
		}
		//如果根结点被调整成了红色，将其改为黑色，并会不影响左右黑结点的个数
		if(_root->_col == RED)
			_root->_col = BLACK;
		return true;
	}

	void InOrder()
	{
		_InOrder(_root);
		cout<<endl;
	}

	/*bool Find(const K& key)
	{

	}*/

private:
	void _Clear(Node* root)
	{
		if(root == NULL)
			return;

		_Clear(root->_left);
		_Clear(root->_right);
		delete root;
		root = NULL;
	}
	
	void _RotateL(Node* parent)//左单旋
    {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
 
        parent->_right = subRL;
        if (subRL != NULL)
            subRL->_parent = parent;
         
        Node* ppNode = parent->_parent;
 
        subR->_left = parent;
        parent->_parent = subR;
 
        if (ppNode == NULL)
            _root = subR;
        else
        {
            if (ppNode->_left == parent)
                ppNode->_left = subR;
            else
                ppNode->_right = subR;
        }
		subR->_parent = ppNode;
    }

    void _RotateR(Node* parent)//右单旋
    {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
 
        parent->_left = subLR;
        if (subLR != NULL)
            subLR->_parent = parent;
 
        Node* ppNode = parent->_parent;
        subL->_right = parent;
        parent->_parent = subL;
 
        if (ppNode == NULL)
            _root = subL;
        else
        {
            if (ppNode->_left == parent)
                ppNode->_left = subL;
            else
                ppNode->_right = subL;
        }
		subL->_parent = ppNode;
	}

	void _InOrder(Node* root)
	{
		if(root != NULL)
		{
			_InOrder(root->_left);
			cout<<root->_key<<" ";
			_InOrder(root->_right);
		}
	}

private:
	Node* _root;
};


void Test()
{
	int arr[] = {3, 4, 6, 1, 7, 2, 8};
	RBTree<int, int> t;
	for(int i = 0; i < sizeof(arr)/sizeof(arr[0]); ++i)
		t.Insert(arr[i], i);

	t.InOrder();
}