C++基础（16）——用红黑树封装出map和set

西阳未落

1711人浏览 · 2025-09-18 12:14:19

西阳未落 · 2025-09-18 12:14:19 发布

红黑树的代码

我们在数据结构专栏中详细的说明了这个结构，感兴趣的友友可跳转到红黑树我们这里直接贴出源码：

#pragma once

enum Colour
{
	RED,
	BLACK
};

template<class K, class V>
struct RBTreeNode
{
	// 这里更新控制平衡也要加入parent指针
	pair<K, V> _kv;
	RBTreeNode<K, V>* _left;
	RBTreeNode<K, V>* _right;
	RBTreeNode<K, V>* _parent;
	Colour _col;

	RBTreeNode(const pair<K, V>& kv)
		:_kv(kv)
		, _left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
	{}
};

template<class K, class V>
class RBTree
{
	typedef RBTreeNode<K, V> Node;
public:
	bool Insert(const pair<K, V>& kv)
	{
		if (_root == nullptr)
		{
			_root = new Node(kv);
			_root->_col = BLACK;

			return true;
		}

		Node* parent = nullptr;
		Node* cur = _root;
		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);
		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 (parent == grandfather->_left)
			{
				//   g
				// p   u
				Node* uncle = grandfather->_right;
				if (uncle && uncle->_col == RED)
				{
					// 变色
					parent->_col = uncle->_col = BLACK;
					grandfather->_col = RED;

					// 继续往上处理
					cur = grandfather;
					parent = cur->_parent;
				}
				else
				{
					if (cur == parent->_left)
					{
						//     g
						//   p    u
						// c
						RotateR(grandfather);
						parent->_col = BLACK;
						grandfather->_col = RED;
					}
					else
					{
						//      g
						//   p    u
						//     c
						RotateL(parent);
						RotateR(grandfather);
						cur->_col = BLACK;
						grandfather->_col = RED;
					}

					break;
				}
			}
			else
			{
				//   g
				// u   p
				Node* uncle = grandfather->_left;
				// 叔叔存在且为红，-》变色即可
				if (uncle && uncle->_col == RED)
				{
					parent->_col = uncle->_col = BLACK;
					grandfather->_col = RED;

					// 继续往上处理
					cur = grandfather;
					parent = cur->_parent;
				}
				else // 叔叔不存在，或者存在且为黑
				{
					// 情况二：叔叔不存在或者存在且为黑
					// 旋转+变色
					//   g
					// u   p
					//       c
					if (cur == parent->_right)
					{
						RotateL(grandfather);
						parent->_col = BLACK;
						grandfather->_col = RED;
					}
					else
					{
						RotateR(parent);
						RotateL(grandfather);
						cur->_col = BLACK;
						grandfather->_col = RED;
					}

					break;
				}
			}
		}

		_root->_col = BLACK;

		return true;
	}

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

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

		Node* pParent = parent->_parent;

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

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

			subL->_parent = pParent;
		}
	}

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

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

	void InOrder()
	{
		_InOrder(_root);
		cout << endl;
	}

	int Height()
	{
		return _Height(_root);
	}

	int Size()
	{
		return _Size(_root);
	}

	Node* Find(const K& key)
	{
		Node* cur = _root;
		while (cur)
		{
			if (cur->_kv.first < key)
			{
				cur = cur->_right;
			}
			else if (cur->_kv.first > key)
			{
				cur = cur->_left;
			}
			else
			{
				return cur;
			}
		}

		return nullptr;
	}

	bool IsBalance()
	{
		if (_root == nullptr)
			return true;

		if (_root->_col == RED)
			return false;

		// 参考值
		int refNum = 0;
		Node* cur = _root;
		while (cur)
		{
			if (cur->_col == BLACK)
			{
				++refNum;
			}
			cur = cur->_left;
		}

		return Check(_root, 0, refNum);
	}

private:

	bool Check(Node* root, int blackNum, const int refNum)
	{
		if (root == nullptr)
		{
			// 前序遍历走到空时，意味着一条路径走完了
			//cout << blackNum << endl;
			if (refNum != blackNum)
			{
				cout << "存在黑色结点的数量不相等的路径" << endl;
				return false;
			}
			return true;
		}

		// 检查孩子不太方便，因为孩子有两个，且不一定存在，反过来检查父亲就方便多了
		if (root->_col == RED && root->_parent->_col == RED)
		{
			cout << root->_kv.first << "存在连续的红色结点" << endl;
			return false;
		}

		if (root->_col == BLACK)
		{
			blackNum++;
		}

		return Check(root->_left, blackNum, refNum)
			&& Check(root->_right, blackNum, refNum);
	}

	void _InOrder(Node* root)
	{
		if (root == nullptr)
		{
			return;
		}

		_InOrder(root->_left);
		cout << root->_kv.first << ":" << root->_kv.second << endl;
		_InOrder(root->_right);
	}

	int _Height(Node* root)
	{
		if (root == nullptr)
			return 0;
		int leftHeight = _Height(root->_left);
		int rightHeight = _Height(root->_right);
		return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
	}

	int _Size(Node* root)
	{
		if (root == nullptr)
			return 0;

		return _Size(root->_left) + _Size(root->_right) + 1;
	}

private:
	Node* _root = nullptr;
};

模板参数

我们根据上面那节可以知道我们的set的实现是只需要K的，而我们的map是需要KV的，那么我们如何两个都是在一棵红黑树中？

我们这里可以重新设计红黑树的模板参数：

template<class K, class T>

在我们的set中，我们的T就是K了：

template<class K>
class set {
public:
    // ...
private:
    RBTree<K, K> _t;
};

我们的map就是正常的传入了，也就是第一个参数给成K，第二个参数给成键值对即可：

template<class K, class V>
class map {
public:
    // ...
private:
    RBTree<K, pair<K, V>> _t;
}

敲黑板：

我们这里不能将我们的第一个参数省略掉，虽然我们的set可以只有一个参数，但是我们的map有的接口是要给出键值的，比如我们的find和我们的erase。

节点中存储的数据

我们的红黑树的模板参数变成变成了K和T，那么我们的红黑树中存储什么：
我们上面也说了，我们根据容器的不同，底层存储的数据也是不同的：

对于set，我们K和T都是存储的Key值。
对于map，我们K代表的就是Key，T代表的就是Key和Value所构成的键值对。

这样之后对于我们set底层存储K和T都是一样的，单数对于我们就只能是存储T了，所以我们考量一下于是就在红黑树的节点中只是存储T了。这样做之后我们的当我们我们的容器是set在红黑树就是存Key，当我们的容器是map在红黑树存的就是<Key, Value>的键值对了。

于是就有了下面这个代码：

template<class T>
struct RBTreeNode
{
	// 这里更新控制平衡也要加入parent指针
	T _data;
	RBTreeNode<T>* _left;
	RBTreeNode<T>* _right;
	RBTreeNode<T>* _parent;
	Colour _col;

	RBTreeNode(const T& data)
		:_data(data)
		, _left(nullptr)
		, _right(nullptr)
		, _parent(nullptr)
	{}
};

模板参数中的仿函数的增加

我们现在有了两种情况了，一种是我们的节点的T存储的是Key，也可能是<Key, Value>键值对，那么我们应该怎么随时获取T中的Key值呢？

我们之所以要随时获取这个Key值是因为我们要堆键的值进行比较操作，我们这个时候会想到可以让我们的map给底层红黑树提供一个仿函数，也就是获取T中的Key值，而对于我们的set而言它的的T本身就是Key值。

我们这里实现仿函数就是在类中实现一个operator()，使得这个类有一个类似函数的行为。

template<class K, class V>
class map
{
    struct MapKeyOfT
    {
        const K& operator()(const pair<K, V>& kv)
        {
            return kv.first;
        }
    };
public:
    // ...
private:
    RBTree<K, pair<K, V>, MapKeyOfT> _t;
};

但是对于我们的底层红黑树来说，它是不知道我们上层的容器是map还是set的，所以我们在进行节点比较的时候都是通过传入的仿函数来获取Key值的。

所以我们的set容器也是要向底层的红黑树传入一个仿函数的，虽然这个仿函数看起来用处不大：

template<class K>
class set
{
    struct SetKeyOfT
    {
        const K& operator()(const K& key)
        {
            return key;
        }
    };
public:
    // ...
private:
    RBTree<K, K, SetKeyOfT> _t;
};

这样之后我们在底层比较的时候都是先通过我们的仿函数获取Key再进行比较，下面我们以我们的查找函数为例：

Node* Find(const K& key) {
    Node* cur = _root;
    while(cur) {
        if(cur->_kv.first < key) {
            cur = cur->_left;
        }else if(cur->_kv.first > key) {
            cur = cur->_right;
        }else {
            return cur;
        }
    }
    return nullptr;
}

正向迭代器的实现

我们实现的正向迭代器实际上就是堆节点的一层封装，所以我们的正向迭代器的实现的成员就是一个节点类型的指针：

template<class T, class Ref, class Ptr>
struct RBTreeIterator
{
	typedef RBTreeNode<T> Node;
	typedef RBTreeIterator<T, Ref, Ptr> Self;

	Node* _node;
}

构造函数

RBTreeIterator(Node* node, Node* root)
		:_node(node)
		,_root(root)
	{}

解引用操作

Ref operator*()
{
    return _node->_data;
}

->操作

直接就是返回节点数据的指针即可

Ptr operator->()
{
    return &_node->_data;
}

实现==和!=操作

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

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

我们实现正向迭代器最难的地方实际上是自增和自减操作

我们实现的正向迭代器进行的自增操作实际上就是中序遍历序列这种当前节点的下一个节点。

逻辑如下：

1、当前的节点右子树不是空，那么自增操作后应该是我们右子树的最左节点。
2、当前节点右子树是空的的，那么自增就是在该节点的祖先节点中找孩子是祖先左的节点。

代码如下：

Self operator++()
{
    if (_node->_right)
    {
        // 右不为空，中序下一个访问的节点是右子树的最左(最小)节点
        Node* min = _node->_right;
        while (min->_left)
        {
            min = min->_left;
        }

        _node = min;
    }
    else
    {
        // 右为空，祖先里面孩子是父亲左的那个祖先
        Node* cur = _node;
        Node* parent = cur->_parent;
        while (parent && cur == parent->_right)
        {
            cur = parent;
            parent = cur->_parent;
        }

        _node = parent;
    }

    return *this;
}

实现红黑的正向迭代你器的自减操作，就是找到在中序遍历的序列中找当前节点的前一个节点。

具体逻辑如下：

1、当前节点的左子树不为空，那么自减后应该找到其左子树的最右节点。
2、当前节点的左子树为空，那么自减后应该在该节点的祖先节点中，找到孩子是父亲右的祖先。
3、如果是到了空节点，我们就找到中序遍历序列的最后一个节点，也就是最右的节点。

代码如下：

Self operator--()
{
    if (_node == nullptr)  // --end()
    {
        // --end()，特殊处理，走到中序最后一个结点，整棵树的最右结点
        Node* rightMost = _root;
        while (rightMost && rightMost->_right)
        {
            rightMost = rightMost->_right;
        }
        _node = rightMost;
    }
    else if (_node->_left)
    {
        // 左子树不为空，中序左子树最后一个
        Node* rightMost = _node->_left;
        while (rightMost->_right)
        {
            rightMost = rightMost->_right;
        }
        _node = rightMost;
    }
    else
    {
        // 孩子是父亲右的那个祖先
        Node* cur = _node;
        Node* parent = cur->_parent;
        while (parent && cur == parent->_left)
        {
            cur = parent;
            parent = cur->_parent;
        }
        _node = parent;
    }

    return *this;
}

我们实现完了上面的几个功能之后，就需要实现我们迭代器中的两个标志性迭代器了begin和end：

begin函数返回的是中序遍历序列的第一个节点的迭代器。
end函数返回的就是中序遍历序列的最后一个节点下一个位置的迭代器，我们这里直接用空指针来构造。

代码如下：

template<class K, class T, class KeyOfT>
class RBTree
{
	typedef RBTreeNode<T> Node;
public:
	typedef RBTreeIterator<T, T&, T*> Iterator;
	typedef RBTreeIterator<T, const T&, const T*> ConstIterator;

	Iterator Begin()
	{
		Node* cur = _root;
		while (cur && cur->_left)
		{
			cur = cur->_left;
		}

		return Iterator(cur, _root);
	}

	Iterator End()
	{
		return Iterator(nullptr, _root);
	}
private:
	Node* _root = nullptr;
};

我们在C++STL中，底层红黑树的正向迭代器的实现。

我们首先要清楚，我们这里实现的效果是有明显缺陷的，理论上我们的在end()为位置的正向迭代器--之后就会得到最后一个位置的正向迭代器，但是我们实现end()的时候，是直接返回nullptr构造的正向迭代器的，所以无法完成上述操作。

底层实现的逻辑图如下：

我们的C++STL库里面实现的红黑树的时候，在红黑树的根节点的位置增加了一个头节点，这个节点的左指针指向我们的红黑树的最左节点，右指针指向我们红黑树的最右节点，父指针指向我们的红黑树的根节点。这个时候我们实现begin()就是返回头节点的左的迭代器，返回end()就是头节点构造的迭代器，end()--的时候就是返回头节点的右的迭代器。

但是我们这个结构的实现还是比较复杂的，因为如果是插入最左节点的左或是插入最右节点的右，这个时候我们需要更新出来头节点的指向，这里就不实现了。