C++ 泛型编程 实现红黑树RBTree

Posted 2021-05-29 Wecccccccc

代码如下:

#include <iostream>
#include <ctime>
using namespace std;

enum COLOR
{
	BLACK,RED
};

template<typename T>
struct RBTreeNode
{
	RBTreeNode<T> * _parent;
	RBTreeNode<T> * _left;
	RBTreeNode<T> * _right;
	T _val;
	COLOR _color;

	RBTreeNode(const T & val = T()):_parent(nullptr),_left(nullptr),_right(nullptr),_val(val),_color(RED){}

};

template<typename T>
class RBTree
{
public:
	typedef RBTreeNode<T> Node;

	RBTree():_root(nullptr){}

	bool insert(const T & val)
	{
		if (_root == nullptr) {
			_root = new Node(val);
			_root->_color = BLACK;
			return true;
		}

		Node *cur = _root;
		Node *parent = nullptr;

		//1.寻找要插入的结点的位置
		while (cur)
		{
			parent = cur;
			if (cur->_val == val) return false;
			else if (cur->_val > val) cur = cur->_left;
			else cur = cur->_right;
		}

		//2.创建结点
		cur = new Node(val);
		if (parent->_val > cur->_val) parent->_left = cur;
		else parent->_right = cur;
		cur->_parent = parent;

		//3.颜色的修改或者结构的调整
		while (cur != _root && cur->_parent->_color == RED)//不为根且存在连续红色，则需要调整
		{
			parent = cur->_parent;
			Node *gfather = parent->_parent;

			if (gfather->_left == parent)
			{
				Node *uncle = gfather->_right;
				//情况1.uncle存在且为红
				if (uncle && uncle->_color == RED)
				{
					parent->_color = uncle->_color = BLACK;
					gfather->_color = RED;
					//向上追溯
					cur = gfather;
				}
				else
				{
					if (parent->_right == cur)
					{
						RotateL(parent);
						swap(cur, parent);
					}
					//情况2.uncle不存在或者uncle为黑
					RotateR(gfather);
					parent->_color = BLACK;
					gfather->_color = RED;
					break;
				}
			}
			else
			{
				Node *uncle = gfather->_left;
				if (uncle && uncle->_color == RED)
				{
					parent->_color = uncle->_color = BLACK;
					gfather->_color = RED;
					//向上追溯
					cur = gfather;
				}
				else
				{
					if (parent->_left == cur)
					{
						RotateR(parent);
						swap(cur, parent);
					}

					RotateL(gfather);
					parent->_color = BLACK;
					gfather->_color = RED;
					break;
				}
			}
		}
		//根结点都为黑色
		_root->_color = BLACK;
		return true;
	}

	void RotateL(Node *parent)
	{
		Node *subR = parent->_right;
		Node *subRL = subR->_left;

		parent->_right = subRL;
		if (subRL) subRL->_parent = parent;

		if (parent == _root)
		{
			_root = subR;
			subR->_parent = nullptr;
		}
		else
		{
			Node *gfather = parent->_parent;
			if (gfather->_left == parent) gfather->_left = subR;
			else gfather->_right = subR;
			subR->_parent = gfather;
		}
		subR->_left = parent;
		parent->_parent = subR;
	}


	void RotateR(Node * parent)
	{
		Node *subL = parent->_left;
		Node *subLR = subL->_right;

		parent->_left = subLR;
		if (subLR) subLR->_parent = parent;

		if (parent == _root)
		{
			_root = subL;
			subL->_parent = nullptr;
		}
		else
		{
			Node *gfather = parent->_parent;
			if (gfather->_left == parent) gfather->_left = subL;
			else gfather->_right = subL;
			subL->_parent = gfather;
		}

		subL->_right = parent;
		parent->_parent = subL;
	}

	void inorder()
	{
		_inorder(_root);
		cout << endl;
	}
//红黑树的性质:
//	1.每个结点的颜色只有红色和黑色，且根结点必须为黑色
//	2.不能存在连续的红色结点
//	3.对于每个结点来说，从该结点到其所有后代结点的简单路径上，经过的黑色结点的数目应该是相同的

	bool isRBTree()
	{
		if (_root == nullptr) return true;//空树也属于红黑树
		if (_root->_color == RED) return false;//不满足性质2：根结点必须为黑色

		int bCount = 0;//从某条路径上的黑色结点数
		Node *cur = _root;
		while (cur)
		{
			if (cur->_color == BLACK) ++bCount;
			cur = cur->_left;
		}
		int pathCount = 0;
		return _isRBTree(_root, bCount, pathCount);
	}

	bool _isRBTree(Node *root, const int bCount, int pathCount)
	{
		if (root == nullptr)//路径走完，判断黑色结点个数
		{
			if (pathCount == bCount) return true;
			else return false;//不满足性质3：每一条路径上的黑色结点个数相同
		}

		if (root->_color == BLACK) ++pathCount;
		Node *parent = root->_parent;
		if (parent && parent->_color == RED && root->_color == RED) return false;//不满足性质2：不能存在连续的红色结点

		return _isRBTree(root->_left, bCount, pathCount) && _isRBTree(root->_right, bCount, pathCount);
	}

private:
	Node *_root;
	void _inorder(Node *root)
	{
		if (root)
		{
			_inorder(root->_left);
			cout << root->_val << " ";
			_inorder(root->_right);
		}
	}
};

int main()
{
	RBTree<int> rbt;
	srand(time(nullptr));
	int n;
	cout << "num = " << endl;
	cin >> n;
	for (int i = 0; i < n; i++)
	{
		rbt.insert(rand());
	}
	rbt.inorder();
	if (rbt.isRBTree()) cout << "yes" << endl;
	else cout << "no" << endl;
	return 0;
}

测试结果: