实现C++中的红黑树

实现C++中的红黑树

RBTree.h

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

enum Color { RED, BLACK };
template<class T>
struct RBTreeNode //红黑树节点
{
	RBTreeNode(const T& x = T(), Color c = RED)
		:left(nullptr)
		, right(nullptr)
		, parent(nullptr)
		, data(x)
		, color(c)
	{}

	RBTreeNode<T>* left;
	RBTreeNode<T>* right;
	RBTreeNode<T>* parent;
	T data;
	Color color;
};

template<class T>
class RBTree
{
	typedef RBTreeNode<T> Node;
public:
	RBTree()
	{
		head = new Node();
		head->left = head;
		head->right = head;
	}
	bool Insert(const T& data)
	{
		//如果树为空
		Node*& root = GetRoot();
		if (nullptr == root)
		{
			root = new Node(data, BLACK);
			root->parent = head;
			head->left = root;
			head->right = root;
			return true;
		}

		//如果树非空先按照二叉搜索树的规则查找要插入的节点位置
		Node* cur = root;
		Node* parent = nullptr;
		while (cur)
		{
			parent = cur;//保存待插入元素的双亲
			if (data < cur->data)
				cur = cur->left;
			else if (data > cur->data)
				cur = cur->right;
			else
				return false;
		}

		//找到位置后按二叉搜索树规则插入节点
		cur = new Node(data);
		if (data < parent->data)
			parent->left = cur;
		else
			parent->right = cur;

		cur->parent = parent;

		// 上述结构满足红黑树的性质则不需要以下调整，若违反红黑树的性质则需要调整，分两侧处理，每侧三种情况
		while (parent != head && RED == parent->color)
		{
			// 违反了性质三，此时需要对红黑树进行调整
			Node* grandfather = parent->parent;
			if (parent == grandfather->left)//双亲在祖父节点的左子树上
			{
				Node* uncle = grandfather->right;
					//情况1：cur为红，p为红，g为黑，u存在且为红
					if (uncle && RED == uncle->color)
					{
						parent->color = BLACK;
						uncle->color = BLACK;
						grandfather->color = RED;
						cur = grandfather;
						parent = cur->parent;
					}
					else
					{
						//情况二: （cur在parent->left）cur为红，p为红，g为黑，u不存在 / u为黑			
						//情况三: （cur在parent->right）cur为红，p为红，g为黑，u不存在 / u为黑
						//情况二和情况三可以合并处理
						if (cur == parent->right)//以parent为根左单旋
						{
							RotateLeft(parent);
							swap(parent, cur);
						}
						//转换成了情况二
						parent->color = BLACK;
						grandfather->color = RED;
						RotateRight(grandfather);
					}
			}
			else//双亲在祖父节点的右子树上
			{
				Node* uncle = grandfather->left;
					//情况1：（不变）
					if (uncle && RED == uncle->color)
					{
						parent->color = BLACK;
						uncle->color = BLACK;
						grandfather->color = RED;
						cur = grandfather;
						parent = cur->parent;
					}
					else
					{
						//情况二: （cur在parent->right）		
						//情况三: （cur在parent->left）
						if (cur == parent->left)//以parent为根右单旋
						{
							RotateRight(parent);
							swap(parent, cur);
						}
						//转换成了情况二
						parent->color = BLACK;
						grandfather->color = RED;
						RotateLeft(grandfather);
					}
			}
		}
		head->left = MostLeft();//头节点的左指向最左边的节点（最小的）
		head->right = MostRight();//头节点的右指向最右边的节点（最大的）
		root->color = BLACK;//控制根节点的颜色为黑，防止树被污染
		return true;
	}
	void InOrder()
	{
		InOrder(head->parent);
		cout << endl;
	}
	bool IsValidRBTree()
	{
		Node* pRoot = GetRoot();

		// 空树也是红黑树
		if (nullptr == pRoot)
			return true;

		// 检测根节点是否满足情况
		if (BLACK != pRoot->color)
		{
			cout << "违反红黑树性质二：根节点必须为黑色" << endl;
			return false;
		}

		// 获取任意一条路径中黑色节点的个数
		size_t blackCount = 0;
		Node* pCur = pRoot;
		while (pCur)
		{
			if (BLACK == pCur->color)
				blackCount++;
			pCur = pCur->left;
		}
		// 检测是否满足红黑树的性质，k用来记录路径中黑色节点的个数
		size_t k = 0;
		return _IsValidRBTree(pRoot, k, blackCount);
	}

private:
	Node*& GetRoot()
	{
		return head->parent;
	}
	Node* MostLeft()
	{
		Node* cur = GetRoot();
		if (nullptr == cur)
			return head;
		else
			while (cur->left)
				cur = cur->left;
		return cur;
	}
	Node* MostRight()
	{
		Node* cur = GetRoot();
		if (nullptr == cur)
			return head;
		else
			while (cur->right)
				cur = cur->right;
		return cur;
	}
	void RotateLeft(Node* parent)
	{
		Node* subR = parent->right;
		Node* subRL = subR->left;

		parent->right = subRL;
		if (subRL)
			subRL->parent = parent;

		subR->left = parent;
		//保存原来parent的双亲，后续需要更新subR的双亲
		Node* Pparent = parent->parent;
		parent->parent = subR;
		subR->parent = Pparent;

		if (Pparent == head)
		{
			//旋转之前parent为根节点，旋转之后parent被subL取代
			head->parent = subR;
		}
		else
		{
			// 旋转之前parent的双亲是存在的
			if (parent == Pparent->left)//Pparent的left还没更新还是subR
				Pparent->left = subR;
			else//Pparent的right还没更新还是subR
				Pparent->right = subR;
		}
	}
	void RotateRight(Node* parent)
	{
		Node* subL = parent->left;
		Node* subLR = subL->right;

		parent->left = subLR;
		if (subLR)
			subLR->parent = parent;

		subL->right = parent;
		//保存原来parent的双亲，后续需要更新subL的双亲
		Node* Pparent = parent->parent;
		parent->parent = subL;
		subL->parent = Pparent;

		if (Pparent == head)
		{
			//旋转之前parent为根节点，旋转之后parent被subL取代
			head->parent = subL;
		}
		else
		{
			// 旋转之前parent的双亲是存在的
			if (parent == Pparent->left)
				Pparent->left = subL;
			else
				Pparent->right = subL;
		}
	}
	void InOrder(Node* root)
	{
		if (root)
		{
			InOrder(root->left);
			cout << root->data << " ";
			InOrder(root->right);
		}
	}
	bool _IsValidRBTree(Node* pRoot, size_t k, const size_t blackCount)
	{
		//走到null之后，判断k和black是否相等
		if (nullptr == pRoot)
		{
			if (k != blackCount)
			{
				cout << "违反性质四：每条路径中黑色节点的个数必须相同" << endl;
				return false;
			}
			return true;
		}
		// 统计黑色节点的个数
		if (BLACK == pRoot->color)
			k++;

		// 检测当前节点与其双亲是否都为红色
		Node* pParent = pRoot->parent;
		if (pParent != head && RED == pParent->color && RED == pRoot->color)
		{
			cout << "违反性质三：没有连在一起的红色节点" << endl;
			return false;
		}

		return _IsValidRBTree(pRoot->left, k, blackCount) &&
			_IsValidRBTree(pRoot->right, k, blackCount);
	}
private:
	Node* head;
};

void TestRBTree()
{
	int a[] = { 5, 3, 4, 1, 7, 8, 2, 6 , 0, 9 };

	RBTree<int> t;
	for (auto e : a)
		t.Insert(e);

	t.InOrder();

	if (t.IsValidRBTree())
	{
		cout << "t is RBTree!!!" << endl;
	}
	else
	{
		cout << "t is not RBTree!!!" << endl;
	}
}

Test.cpp

#include "RBTree.h"
using namespace std;

int main()
{
	TestRBTree();
	return 0;
}

运行结果