【数据结构】带迭代器的红黑树

在上一篇博客中，我们简单的介绍了红黑树及其插入操作，下面我们将给红黑树封装一个迭代器。为方便操作，我们需要添加一个头结点，并设置其颜色为红色，以便于根结点区分开；这个根结点的左指向最小的结点，右指向树中值最大的结点，其双亲结点指向根结点，根结点的双亲域也指向头，如下图所示：

#pragma once
#include <iostream>
using namespace std;

enum COLOR { RED, BLACK };
template <class K, class V>
struct RBNode
{
	RBNode(const K& key = K(), const V& value = V(), const COLOR color = RED)
	: _pLeft(NULL)
	, _pRight(NULL)
	, _pParent(NULL)
	, _key(key)
	, _value(value)
	, _color(color)
	{}
	RBNode<K, V>* _pLeft;
	RBNode<K, V>* _pRight;
	RBNode<K, V>* _pParent;
	K _key;
	V _value;
	COLOR _color;
};

template <class K, class V, class Ref, class Ptr>
struct RBTreeIterator
{
private:
	typedef RBNode<K, V> Node;
	typedef Node* pNode;
	typedef RBTreeIterator<K, V, Ref, Ptr> Self;
public:
	RBTreeIterator()
		: _pNode(NULL)
	{}

	RBTreeIterator(pNode pnode)
		: _pNode(pnode)
	{}

	RBTreeIterator(Self& s)
		: _pNode(s._pNode)
	{}

	bool operator==(const Self& s)
	{
		return _pNode == s._pNode;
	}

	bool operator!=(const Self& s)
	{
		return _pNode != s._pNode;
	}

	Ptr operator->()
	{
		return &(_pNode->_key);
	}

	Ref operator*()
	{
		return _pNode->_key;
	}

	Self& operator++()
	{
		Increasement();
		return *this;
	}

	Self operator++(int)
	{
		Self tmp(*this);
		Increasement();
		return tmp;
	}

	Self& operator--()
	{
		Decreasement();
		return *this;
	}

	Self operator--(int)
	{
		Self tmp(*this);
		Decreasement();
		return tmp;
	}

private:
	void Increasement()
	{
		if (_pNode->_pRight)
		{
			_pNode = _pNode->_pRight;
			while (_pNode->_pLeft)
				_pNode = _pNode->_pLeft;
		}
		else
		{
			pNode pParent = _pNode->_pParent;
			while (_pNode == pParent->_pRight)
			{
				_pNode = pParent;
				pParent = _pNode->_pParent;
			}
			if (_pNode->_pRight != pParent)
				_pNode = pParent;
		}
	}

	void Decreasement()
	{
		if (_pNode->_pParent->_pParent==_pNode && _pNode->_color==RED)
			_pNode = _pNode->_pRight;
		else if (_pNode->_pLeft)
		{
			_pNode = _pNode->_pLeft;
			while (_pNode->_pRight)
				_pNode = _pNode->_pRight;
		}
		else
		{
			pNode pParent = _pNode->_pParent;
			while (_pNode == pParent->_pLeft)
			{
				_pNode = pParent;
				pParent = _pNode->_pParent;
			}
			_pNode = pParent;
		}
	}
private:
	pNode _pNode;
};

template <class K, class V>
class RBTree
{
	typedef RBNode<K, V> Node;
	typedef Node* pNode;
public:
	typedef RBTreeIterator<K, V, K&, K*> Iterator;
	RBTree()
	{
		_pHead = new Node();
		_pHead->_color = RED;
		_pHead->_pLeft = _pHead;
		_pHead->_pRight = _pHead;
		_pHead->_pParent = NULL;
	}

	pair<bool, pNode> Insert(const K& key, const V& value)
	{
		pNode& _pRoot = GetRoot();
		if (NULL == _pRoot)
		{
			_pRoot = new Node(key, value);
			_pRoot->_color = BLACK;
			_pRoot->_pParent = _pHead;
			_pHead->_pLeft = _pRoot;
			_pHead->_pRight = _pRoot;
			_pHead->_pParent = _pRoot;
			_count++;
			return pair<bool, pNode>(true, _pRoot);
		}
		pNode pCur = _pRoot;
		pNode pParent = NULL;
		while (pCur)
		{
			pParent = pCur;
			if (key < pCur->_key)
				pCur = pCur->_pLeft;
			else
				pCur = pCur->_pRight;
		}
		pCur = new Node(key, value);
		if (key < pParent->_key)
			pParent->_pLeft = pCur;
		else if (key > pParent->_key)
			pParent->_pRight = pCur;
		else
			return make_pair(false, (pNode)NULL);
		pCur->_pParent = pParent;
		//调节颜色
		while (pCur!=_pRoot && pParent->_color==RED)
		{
			pNode grandFather = pParent->_pParent;
			if (pParent == grandFather->_pLeft)
			{
				pNode uncle = grandFather->_pRight;
				if (uncle && uncle->_color == RED)//情况三
				{
					pParent->_color = BLACK;
					uncle->_color = BLACK;
					grandFather->_color = RED;
					pCur = grandFather;
					pParent = pCur->_pParent;
				}
				else
				{
					if (pCur == pParent->_pRight)//情况五
					{
						RotateL(pParent);
						swap(pParent, pCur);
					}
					pParent->_color = BLACK;//情况四
					grandFather->_color = RED;
					RotateR(grandFather);
				}
			}
			else
			{
				pNode uncle = grandFather->_pLeft;
				if (uncle && uncle->_color == RED)
				{
					pParent->_color = BLACK;
					uncle->_color = BLACK;
					grandFather->_color = RED;
					pCur = grandFather;
					pParent = pCur->_pParent;
				}
				else
				{
					if (pCur == pParent->_pLeft)
					{
						RotateR(pParent);
						swap(pParent, pCur);
					}
					pParent->_color = BLACK;
					grandFather->_color = RED;
					RotateL(grandFather);
				}
			}
		}
		_pRoot->_color = BLACK;
		_pHead->_pLeft = MinNode();
		_pHead->_pRight = MaxNode();
		_count++;
		return make_pair(true, pCur);
	}

	Iterator Begin()
	{
		return Iterator(_pHead->_pLeft);
	}

	Iterator End()
	{
		return Iterator(_pHead);
	}

	int Size()
	{
		return _count;
	}
private:
	pNode& GetRoot()
	{
		return _pHead->_pParent;
	}

	pNode MinNode()
	{
		if (NULL == _pHead->_pParent)
			return NULL;
		pNode pCur = _pHead->_pParent;
		while (pCur->_pLeft)
			pCur = pCur->_pLeft;
		return pCur;
	}

	pNode MaxNode()
	{
		if (NULL == _pHead->_pParent)
			return NULL;
		pNode pCur = _pHead->_pParent;
		while (pCur->_pRight)
			pCur = pCur->_pRight;
		return pCur;
	}
	void RotateL(pNode pParent)
	{
		pNode& _pRoot = GetRoot();
		pNode pSubR = pParent->_pRight;
		pNode pSubRL = pSubR->_pLeft;

		pParent->_pRight = pSubRL;
		if (pSubRL)
			pSubRL->_pParent = pParent;
		pNode pPParent = pParent->_pParent;
		pSubR->_pLeft = pParent;
		pParent->_pParent = pSubR;
		pSubR->_pParent = pPParent;
		if (pParent == _pRoot)
			_pRoot = pSubR;
		else
		{
			if (pParent == pPParent->_pLeft)
				pPParent->_pLeft = pSubR;
			else
				pPParent->_pRight = pSubR;
		}
	}

	void RotateR(pNode pParent)
	{
		pNode& _pRoot = GetRoot();
		pNode pSubL = pParent->_pLeft;
		pNode pSubLR = pSubL->_pRight;

		pParent->_pLeft = pSubLR;
		if (pSubLR)
			pSubLR->_pParent = pParent;
		pNode pPParent = pParent->_pParent;
		pSubL->_pRight = pParent;
		pParent->_pParent = pSubL;
		pSubL->_pParent = pPParent;
		if (pParent == _pRoot)
			_pRoot = pSubL;
		else
		{
			if (pParent == pPParent->_pLeft)
				pPParent->_pLeft = pSubL;
			else
				pPParent->_pRight = pSubL;
		}
	}
private:
	pNode _pHead;
	int _count;
};

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

	RBTree<int, int>::Iterator it = rbt.Begin();
	while (it != rbt.End())
	{
		cout << *it << " ";
		++it;
	}
	cout << endl;
	cout << rbt.Size() << endl;
}

测试结果如下：