【C++篇】平衡二叉搜索树（下篇）：红黑树详解

// 颜色枚举，用于红黑树节点颜色标记
enum Color
{
    BLACK,  // 黑色
    RED     // 红色
};

// 红黑树节点模板类
template<class K, class V>
struct BRTreeNode
{
    BRTreeNode<K, V>* _left;    // 左子节点指针
    BRTreeNode<K, V>* _right;   // 右子节点指针
    BRTreeNode<K, V>* _parent;  // 父节点指针
    pair<K, V> _kv;             // 存储的键值对
    Color _col;                 // 节点颜色（红或黑）

    // 构造函数
    BRTreeNode(const pair<K, V>& kv)
        :_left(nullptr)     // 初始化左子节点为空
        ,_right(nullptr)    // 初始化右子节点为空
        ,_parent(nullptr)   // 初始化父节点为空
        ,_kv(kv)           // 初始化键值对
        ,_col(RED)         // 新节点默认为红色（根据红黑树规则）
    { }
};

// 红黑树模板类
template<class K, class V>
class BRTree
{
    // 类型重定义，简化代码
    typedef BRTreeNode<K, V> Node;
public:
    // 公共接口
    // 通常包括插入、删除、查找等操作
    //…………
    
private:
    Node* _root = nullptr;  // 红黑树的根节点，初始化为空
};

// 在红黑树中插入值为kv的节点
bool Insert(const pair<K, V>& kv)
{
	if (_root == nullptr)
	{
		_root = new Node(kv);
		_root->_col = BLACK;
		return true;
	}

	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;
		}
	}

	cur = new Node(kv);
	//新增在左
	if (parent->_kv.first > kv.first)
	{
		parent->_left = cur;

	}
	else//新增在右
	{
		parent->_right = cur;
	}
	cur->_parent = parent;

	//若插入破坏平衡，调节至平衡
	//…………
	return true;
}

//更新颜色
//检查是否需要更新颜色，若parent为黑则无需更新
while (parent && parent->_col == RED)
{
	Node* grandfather = parent->_parent;
	Node* uncle = nullptr;
	if (grandfather->_left == parent)
	{
		uncle = grandfather->_right;
	}
	else
	{
		uncle = grandfather->_left;
	}

	if (uncle && uncle->_col == RED)//uncle存在且为红——变色
	{
		parent->_col = uncle->_col = BLACK;
		grandfather->_col = RED;

		if (grandfather->_parent)
		{
			if (grandfather->_col == BLACK)
			{
				break;
			}
			else
			{
				cur = grandfather;
				parent = cur->_parent;
			}
		}
		else
		{
			grandfather->_col = BLACK;
			break;
		}
	}
	else//uncle不存在或uncle存在且为黑——旋转+变色
	{
		//判断何种情况，该用何种旋转
		if (grandfather->_right == parent)
		{
			if (cur == parent->_right)
			{
				//  g
				//     p
				//		  c
				RotateL(grandfather);
				parent->_col = BLACK;
				grandfather->_col = RED;
			}
			else
			{
				//  g
				//     p
				//  c
				RotateR(parent);
				RotateL(grandfather);
				cur->_col = BLACK;
				grandfather->_col = RED;
			}
		}
		else
		{
			if (cur == parent->_left)
			{
				//        g
				//     p
				//  c
				RotateR(grandfather);
				parent->_col = BLACK;
				grandfather->_col = RED;
			}
			else
			{
				//    g
				//  p
				//    c
				RotateL(parent);
				RotateR(grandfather);
				cur->_col = BLACK;
				grandfather->_col = RED;
			}
		}
		break;
	}
}
_root->_col = BLACK;

bool CheckCol(Node* root, int BNum, int benchmark)
{
	//走到null之后，判断k和black是否相等
	if (root == nullptr)
	{
		if (BNum != benchmark)
			return false;
		return true;
	}
	
	// 统计黑色节点的个数
	if (root->_col == BLACK)
		++BNum;
		
	//检查但前节点与父节点是否红色节点连续
	if (root->_col == RED && root->_parent && root->_parent->_col == RED)
	{
		cout << root->_kv.first << "出现连续红色节点" << endl;
		return false;
	}

	return CheckCol(root->_left, BNum, benchmark)
		&& CheckCol(root->_right, BNum, benchmark);
}

bool IsBalance(Node* root)
{
	// 空树也是红黑树
	if (root == nullptr)
	{
		return true;
	}
	// 检测根节点是否满足情况
	if (root->_col != BLACK)
	{
		return false;
	}

	//获取任意一条路径中黑色节点的个数作为基准值
	int benchmark = 0;
	Node* cur = root;
	while (cur)
	{
		if (cur->_col == BLACK)
			benchmark++;
		cur = cur->_left;
	}
	// 检测是否满足红黑树的性质，k用来记录路径中黑色节点的个数
	size_t k = 0;
	return CheckCol(root, k, benchmark);
}