红黑树的平衡之舞：数据结构中的优雅艺术

template<class K, class V>
class RBTreeNode {
public:
	RBTreeNode<K, V>* left;
	RBTreeNode<K, V>* right;
	RBTreeNode<K, V>* _parent;

	pair<K, V> _kv; //结点中存的键值对
	color _col;		//结点的颜色

	RBTreeNode(const& pair<K, V> kv = pair<K, V>(), Color color = RED)
		:_kv(kv)
		,left(nullptr)
		,right(nullptr)
		,_parent(nullptr)
		,_col(color)
	{}
	
};

template<class K, class V>
class RBTree {
	typedef RBTreeNode<K, V> Node;
public:
    // 成员函数...    
private:
	Node* _root;
};

// 左单旋
void RotateL(Node* parent) {
	Node* cur = parent->right;
	Node* curleft = cur->left;

	parent->right = curleft;	
	cur->left = parent;

	if (curleft) curleft->_parent = parent;

	Node* ppnode = parent->_parent;
	parent->_parent = cur;

	if (parent == _root) {
		_root = cur;
		cur->_parent = nullptr; 
	}
	else {
		if (ppnode->left == parent) ppnode->left = cur;
		else ppnode->right = cur;

		cur->_parent = ppnode;
	}
}

// 右单旋
void RotateR(Node* parent) {
	Node* cur = parent->left;
	Node* curright = cur->right;

	parent->left = curright;
	cur->right = parent;

	if (curright) curright->_parent = parent;

	Node* ppnode = parent->_parent;
	parent->_parent = cur;

	if (parent == _root) {
		_root = cur;
		cur->_parent = nullptr;
	}
	else {
		if (ppnode->left == parent) ppnode->left = cur;
		else ppnode->right = cur;

		cur->_parent = ppnode;
	}
}

bool insert(const pr& kv) {
	if (!_root) {
		_root = new Node(kv);
		_root->_col = BLACK;
		return true;
	}

	// 定义当前节点cur
	Node* cur = _root;
	Node* parent = nullptr;
	while (cur) {
		if (cur->_kv.first < kv.first) {
			parent = cur;
			cur = cur->right;
		}
		else if (cur->_kv.first > kv.first) {
			parent = cur;
			cur = cur->left;
		}
		else return false; // 不可能出现相等的情况插入
	}

	// 创建新的Node实例cur，并将它插入到适当的位置
	cur = new Node(kv);
	cur->_col = RED;
	if (parent->_kv.first < kv.first) {
		parent->right = cur;
	}
	else {
		parent->left = cur;
	}

	cur->_parent = parent;

	// 更新结点颜色，控制平衡
	while (parent && parent->_col == RED) {
		Node* grandfather = parent->_parent;
		// 父亲是祖父的左孩子
		if (grandfather->left == parent) {
			// 叔叔存在且为红
			if (uncle && uncle->_col == RED) {
				parent->_col = uncle->_col = BLACK;
				grandfather->_col = RED;
				
				//继续向上处理
				cur = grandfather;
				parent = cur->_parent;
			}
			// 叔叔不存在或者存在且为黑
			else {
				if (cur == parent->left) {
					RotateR(grandparent);
					parent->_col = BLACK;
					grandfather->_col = RED;
				}
				else {
					RotateL(cur);
					RotateR(grandfather);
					cur->_col = BLACK;
					grandfather->_col = RED;
				}
			}
		}
		// 父亲是祖父的右孩子
		else {
			// 叔叔存在且为红
			if (uncle && uncle->_col == RED) {
				parent->_col = uncle->_col = BLACK;
				grandfather->_col = RED;

				//继续向上处理
				cur = grandfather;
				parent = cur->_parent;
			}
			// 叔叔不存在或者存在且为黑
			else {
				if (cur == parent->right) {
					RotateL(grandfather);
					parent->_col = BLACK;
					grandfather->_col = RED;
				}
				else {
					RotateR(cur);
					RotateL(grandfather);
					cur->_col = BLACK;
					grandfather->_col = RED;
				}
			}
		}
	}

	return true;
}

// 调用判断平衡的函数
bool IsBalance() {
	return IsBalance(_root);
}

// 判断平衡
bool IsBalance(Node* root) {
	if (!root) return true;
	if (root->_col != BLACK) return false;

	// 随便找一条路径，这里我们选择最左路径
	int benchmark = 0;
	Node* cur = _root;
	while (cur) {
		if (cur->_col == BLACK) {
			++benchmark;
		}
		cur = cur->left;
	}

	return CheckColor(root, 0, benchmark);
}

// 主要检查有无连续的红色结点
bool CheckColor(Node* root, int blacknum, int benchmark) {
	if (!root) {
		if (blacknum != benchmark) {
			return false;
		}
		return true;
	}

	if (root->_col == BLACK) ++blacknum;

	if (root->_col == RED && root->_parent && root->_parent->_col == RED) {
		cout << root->_kv.first << "出现连续红色节点" << endl;
		return false;
	}

	return CheckColor(root->left, blacknum, benchmark) && CheckColor(root->right, blacknum, benchmark);
}

// 调用中序遍历
void InOrder() {
	return InOrder(_root);
}

// 中序遍历
void InOrder(Node* root) {
	if (!root) return;

	InOrder(root->left);
	cout << root->_kv.first << " ";
	InOrder(root->right);
}

#include "RBTree.h"

int main() {
	int a[] = { 1, 5, 7, 16, 25, 36, 44, 55,6, 32, 9 };
	RBTree<int, int> rbt;
	for (auto e : a) {
		rbt.insert(make_pair(e,e));
		rbt.InOrder();
		cout << endl;
		if (rbt.IsBalance()) {
			cout << "插入成功" << endl;
		}
		else {
			cout << "插入失败" << endl;
		}
	}

	return 0;
}