【C++进阶】 红黑树简单模拟实现STL中的map与set

👦个人主页：@Weraphael
?🏻作者简介：目前学习C++和算法
??专栏：C++航路
🐋 希望大家多多支持，咱一起进步！😁
如果文章对你有帮助的话
欢迎 评论💬 点赞👍🏻 收藏 📂 加关注?

---

一、set和map源码剖析（差源码）

源码地址：

我们首先可以观察到，在set和map中包含有如下的头文件

先剧透一下：我们知道，set是key模型，map是key-value模型，按道理来说应该用两颗不同的红黑树分别来封装key和key-value。但其实在stl中，它们使用的是同一颗树。

那么它是如何做到同一颗树既可以存set，也可以存map呢？

我们首先可以打开stl_set.h头文件以及stl_map.h来剖析。

• stl_set.h头文件

在以往的博客中讲过，看源码首先需要看它的成员变量

我们发现：set的底层好像是key-value结构（和我们一开始说的key模型不太一样），不同的是它们都是Key。接下来我们再来看看stl_map.h头文件

• stl_map.h头文件

我们发现：map是key_value结构，但它的value竟然是个pair结构

也就是说

• set是<K,K>模型的树
• map是<K，pair>模型的树

为什么STL要这样设计呢？接下来我们再看看红黑树stl_tree.h的实现

• stl_tree.h头文件

以上源代码的核心是link_type header，其类型是rb_tree_node*，也就是红黑树结点的指针。

接下来我们需要关心这颗树的Value。大家发现没有 红黑树的结点存储的是Value

但是这个Value不是我们前面所说的key-val模型中的val。

这里的Value对于map而言是pair，对于set而言是key。所以说，真正决定红黑树里面存储的是什么，是由第二个模板参数决定的，这也就为什么map和set可以共用一颗树。其实也不是，更严谨一点来说是共用同一个类模板（红黑树）。

我们现在再来观测一下这个红黑树结点里面有什么

我们发现：这里是通过一个继承关系来搞定的。派生类存储的是value，而基类存储的是三叉链的指针和颜色。

那么在这里我们似乎看到了，在这棵树里面，我们好像并不是很需要key-value结构中的key类型的模板参数，那么事实上是如此的吗？其实不是的，这个key还必须得传入，因为会有一些接口需要key。比如对于map而言，find需要通过key来查找。

二、改造红黑树

2.1 红黑树基本结构的改造

在前头说过：真正决定红黑树里面存储的是什么，是由第二个模板参数决定的。

enum Colour
{
	RED,
	BLACK
};

template <class T>
struct RBTreeNode 
{
	// 结点只需要存红黑树第二个模板参数类型的数据
	T _data;
	struct RBTreeNode<T>* _left;
	struct RBTreeNode<T>* _right;
	struct RBTreeNode<T>* _parent;
	Colour _col;

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

template <class K, class T>
class RBTree
{
	// 真正决定红黑树里面存储的是什么
	// 是由第二个模板参数决定的，也就是T
	typedef struct RBTreeNode<T> Node;

public:
	RBTree()
		:_root(nullptr)
	{}

private:
	Node* _root;
}

同时也可以写出map和set的基本结构

【set.h】

#include "RBTree.h"

namespace wj
{
	template<class K>
	class set
	{
	public:

	private:
		RBTree<K, K> _set;
	};
}

【map.h】

#include "RBTree.h"

namespace wj
{
	template<class K, class V>
	class map
	{
	public:

	private:
		RBTree<K, pair<K, V>> _map;
	};
}

2.2 插入操作的改造

首先我们来分析插入的模板参数应该是什么？对于set就是key；对于map则是pair。那么模板参数应该是T。

那现在就会遇到一个非常棘手的问题：插入的位置代码不怎么好写。

原因是：假设是set，data在实例化后，此时是没有问题的；如果是map，data在实例化后，由于是pair，此处应该比较pair的first，而pair的重载和我们期望的不一样，因此会出错。

解决方法是写一个仿函数。首先在map和set里面写两个内部类，这个内部类里面是一个运算符重载，充当仿函数。对于map返回pair的first，对于set返回它本身，也就是key

• 【set.h】

namespace wj
{
	template<class K>
	class set
	{
	public:
		struct SetKeyOfT
		{
			const K& operator()(const K& key)
			{
				return key;
			}
		};

		bool insert(const K& key)
		{
			return _set.Insert(key);
		}

	private:
		RBTree<K, K, SetKeyOfT> _set;
	};
}

• 【map.h】

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

		bool insert(const pair<K, V>& kv)
		{
			return _map.Insert(kv);
		}

	private:
		RBTree<K, pair<K, V>, MapKeyOfT> _map;
	};
}

• 插入部分

// 因为关联式容器中存储的是<key, value>的键值对，因此
// k为key的类型，
// T: 如果是map，则为pair<K, V>; 如果是set，则为k
// KeyOfT: 通过value来获取key的一个仿函数类
template <class K, class T, class KeyOfT>
class RBTree
{
	// 真正决定红黑树里面存储的是什么
	// 是由第二个模板参数决定的，也就是T
	typedef struct RBTreeNode<T> Node;

public:
	bool Insert(const T& data)
	{
		if (_root == NULL)
		{
			_root = new Node(data);
			_root->_col = BLACK; 
			return true;
		}
	
		Node* parent = nullptr;
		Node* cur = _root;
		
		// 仿函数
		KeyOfT kot;
		
		while (cur)
		{
			// 棘手
			if (kot(cur->_data) < kot(data))
			{
				parent = cur;
				cur = cur->_right;
			}
			else if (kot(cur->_data) > kot(data))
			{
				parent = cur;
				cur = cur->_left;
			}
			else 
			{
				return false;
			}
		}
		cur = new Node(data);
		cur->_col = RED; 
		
	
		if (kot(parent->_data) < kot(data))
		{
			parent->_right = cur;
		}
		else
		{
			parent->_left = cur;
		}
		cur->_parent = parent;
	
		while (parent && parent->_col == RED) 
		{
			Node* grandparent = parent->_parent; 
			if (parent == grandparent->_left)
			{
				Node* uncle = grandparent->_right;
	
				if (uncle && uncle->_col == RED)
				{
					parent->_col = uncle->_col = BLACK;
					grandparent->_col = RED;
					cur = grandparent;
					parent = cur->_parent;
				}
				else
				{
					if (cur == parent->_left)
					{
						RotateRight(grandparent);
						parent->_col = BLACK;
						grandparent->_col = RED;
					}
					else
					{
						RotateLeft(parent);
						RotateRight(grandparent);
	
						cur->_col = BLACK;
						grandparent->_col = RED;
					}
					break;
				}
			}
	
			else 
			{
				Node* uncle = grandparent->_left;
				if (uncle && uncle->_col == RED)
				{
					parent->_col = uncle->_col = BLACK;
					grandparent->_col = RED;
					cur = grandparent;
					parent = cur->_parent;
				}
				else
				{
					if (cur == parent->_right)
					{
						RotateLeft(grandparent);
						grandparent->_col = RED;
						parent->_col = BLACK;
					}
					else
					{
						RotateRight(parent);
						RotateLeft(grandparent);
						cur->_col = BLACK;
						grandparent->_col = RED;
					}
					break;
				}
			}
		}
		_root->_col = BLACK; 
		return true;
	}
private:
	Node* _root;
};

2.3 添加查找操作

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

三、迭代器

首先来看看库里的set的迭代器是如何实现的

如上图所示，迭代器同样是依靠着红黑树里面的迭代器去实现的。那么我们可以去stl_tree.h文件看看

我们可以看到，这个迭代器与list的迭代器是比较相似的。

接下来可以简单看看迭代器里的操作

大家注意到没，这里的++和--操作和我们在模拟实现list的时候一样，封装了一个结点的指针，然后重载运算符。但是这是一颗红黑树啊，++和--往哪走？

以下是迭代器遍历的代码块

auto it = a.begin();
while (it != a.end());
{
	cout << *it << ' ';
	++it;
}

首先可以想到：迭代器遍历的时候默认是升序，也就是中序遍历，那么 a.begin返回的一定是树中的最小值。比如一下这幅图

接下来++应该访问结点6。所以，如果右树为空，那么下一个访问的就是孩子是父亲的左孩子的祖先；接下来访问结点7，那么如果右树不为空，访问的是右树的最左结点（即最小结点）。当然了，--操作就是++反向操作，因为++走的是中序遍历（左子树，根，右子树），那么--就是右子树，根，左子树。

template<class T>
struct __TreeIterator
{
	typedef RBTreeNode<T> Node;
	typedef __TreeIterator<T> Self;

	Node* _node;

	__TreeIterator(Node* node)
		:_node(node)
	{}


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

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

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

	bool operator==(const Self& s)
	{
		return _node == s._node;
	}
	
	Self& operator++()
	{
		if (_node->_right)
		{
			// 右树的最左结点（最小结点）
			Node* subLeft = _node->_right;
			while (subLeft->_left)
			{
				subLeft = subLeft->_left;
			}
			_node = subLeft;
		}
		else
		{
			Node* cur = _node;
			Node* parent = cur->_parent;
			// 找孩子是父亲的左孩子的祖先，就是下一个要访问的结点
			while (parent)
			{
				if (cur == parent->_left)
				{
					break;
				}
				else
				{
					cur = cur->_parent;
					parent = parent->_parent;
				}
			}
			_node = parent;
		}
		return *this;
	}

	Self operator--()
	{
		if (_node->_left)
		{
			Node* subRight = _node->_left;
			while (subRight->_right)
			{
				subRight = subRight->_right;
			}
			_node = subRight;
		}
		else
		{
			Node* cur = _node;
			Node* parent = cur->_parent;
			while (parent && cur == parent->_left)
			{
				cur = cur->_parent;
				parent = parent->_parent;
			}
			_node = parent;
		}
		return *this;
	}
};

有了迭代器类，我们可以在红黑树层次去调用这个迭代器类了

然后我们依次到map和set层次去封装

• 【set.h】

namespace wj
{
	template<class K>
	class set
	{
	public:
		struct SetKeyOfT
		{
			const K& operator()(const K& key)
			{
				return key;
			}
		};

		typedef typename RBTree<K, K, SetKeyOfT>::iterator iterator;
		
		iterator begin()
		{
			return _set.begin();
		}

		iterator end()
		{
			return _set.end();
		}

		bool insert(const K& key)
		{
			return _set.Insert(key);
		}

	private:
		RBTree<K, K, SetKeyOfT> _set;
	};
}

• 【map.h】

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

		typedef typename RBTree<K, pair<K, V>, MapKeyOfT>::iterator iterator;
		iterator begin()
		{
			return _map.begin();
		}

		iterator end()
		{
			return _map.end();
		}

		bool insert(const pair<K, V>& kv)
		{
			return _map.Insert(kv);
		}

	private:
		RBTree<K, pair<K, V>, MapKeyOfT> _map;
	};
}

不过在这里我们似乎发现了一个问题，那就是set容器的数据应该是不可以被修改的，但是我们对他进行了修改，而且还可以成功使用。那就说明我们的迭代器还存在bug，因为一旦修改了set里面的数据，那么就意味着这棵树已经乱了。可能不再是一个搜索二叉树了。

而在库里面是这样实现的，对于set，两个迭代器本质都是const迭代器。对于map，它的key应该不可以被修改，这里是通过对pair的第一个参数first进行const限定的，从而锁死了第一个参数。

下面是红黑树中的修改

当然了，咱们自定义类型的迭代器也得修改

然后我们来修改set中的迭代器函数。

需要注意的是：普通迭代器也可以调用const迭代器（库中也是这么实现的）

然后我们来处理map的迭代器问题

处理方式很简单：只需要让pair中的第一个参数给带上const即可。保证key不可以被修改，而value可以修改。

四、[ ] 操作符

map的[]操作主要是依靠insert操作实现的。所以还得将insert在进一步完善

pair<iterator, bool> Insert(const T& data)
{
	if (_root == NULL)
	{
		_root = new Node(data);
		_root->_col = BLACK; 
		return make_pair(iterator(_root), true);
	}

	Node* parent = nullptr;
	Node* cur = _root;
	
	KeyOfT kot;
	
	while (cur)
	{
		// 棘手
		if (kot(cur->_data) < kot(data))
		{
			parent = cur;
			cur = cur->_right;
		}
		else if (kot(cur->_data) > kot(data))
		{
			parent = cur;
			cur = cur->_left;
		}
		else 
		{
			return make_pair(iterator(cur), false);
		}
	}
	cur = new Node(data);
	cur->_col = RED; 
	Node* newnode = cur;

	if (kot(parent->_data) < kot(data))
	{
		parent->_right = cur;
	}
	else
	{
		parent->_left = cur;
	}
	cur->_parent = parent;

	while (parent && parent->_col == RED) 
	{
		Node* grandparent = parent->_parent; 
		if (parent == grandparent->_left)
		{
			Node* uncle = grandparent->_right;

			if (uncle && uncle->_col == RED)
			{
				parent->_col = uncle->_col = BLACK;
				grandparent->_col = RED;
				cur = grandparent;
				parent = cur->_parent;
			}
			else
			{
				if (cur == parent->_left)
				{
					RotateRight(grandparent);
					parent->_col = BLACK;
					grandparent->_col = RED;
				}
				else
				{
					RotateLeft(parent);
					RotateRight(grandparent);

					cur->_col = BLACK;
					grandparent->_col = RED;
				}
				break;
			}
		}

		else 
		{
			Node* uncle = grandparent->_left;
			if (uncle && uncle->_col == RED)
			{
				parent->_col = uncle->_col = BLACK;
				grandparent->_col = RED;
				cur = grandparent;
				parent = cur->_parent;
			}
			else
			{
				if (cur == parent->_right)
				{
					RotateLeft(grandparent);
					grandparent->_col = RED;
					parent->_col = BLACK;
				}
				else
				{
					RotateRight(parent);
					RotateLeft(grandparent);
					cur->_col = BLACK;
					grandparent->_col = RED;
				}
				break;
			}
		}
	}
	_root->_col = BLACK; 
	return make_pair(iterator(newnode), true);
}

然后我们可以直接修改map和set的返回值

但是当我们运行的时候，报错了，我们发现是set中的pair的迭代器出问题了

这时因为set中的迭代器都是const迭代器，而红黑树是一个普通对象，它返回的是一个普通的迭代器。所以我们上面的写法出现了类型不兼容的问题

我们可以参考库里面是这样处理的

我们可以尝试仿照它写

但是还是编译不通过

这里其实涉及到pair的构造函数了

可以看到，pair其实是通过初始化列表完成的。所以说在将first初始化的时候，会调用它的构造函数。去用它的迭代器去构造一个const迭代器。

这里我们需要注意的一个问题是，这个构造如何写？我们参照库里面的代码。可以看到，库里面专门将普通迭代器和const迭代器给再次typedef了，这样的好处就在于，当这个迭代器类是一个const迭代器的时候，那么这里就是构造函数，如果这个迭代器被实例化为了普通迭代器的时候，这里就是一个拷贝构造函数了

于是我们可以仿照库里面的代码

那么最后这个[]操作就很容易实现了

五、本篇源代码以及stl相关源码

Gitte仓库：点击跳转