【C++】封装红黑树实现mymap和myset

1. 源码及框架分析

1.1 源码分析

SGI-STL30版本源代码，map和set的源代码在map/set/stl_map.h/stl_set.h/stl_tree.h等几个头文件

中。

map和set的实现结构框架核心部分截取出来如下：

cpp 复制代码

// set
#ifndef __SGI_STL_INTERNAL_TREE_H
#include <stl_tree.h>
#endif
#include <stl_set.h>
#include <stl_multiset.h>
// map
#ifndef __SGI_STL_INTERNAL_TREE_H
#include <stl_tree.h>
#endif
#include <stl_map.h>
#include <stl_multimap.h>
// stl_set.h
template <class Key, class Compare = less<Key>, class Alloc = alloc>
class set {
public:
	// typedefs:
	typedef Key key_type;
	typedef Key value_type;
private:
	typedef rb_tree<key_type, value_type,
		identity<value_type>, key_compare, Alloc> rep_type;
	rep_type t; // red-black tree representing set
};
// stl_map.h
template <class Key, class T, class Compare = less<Key>, class Alloc = alloc>
class map {
public:
	// typedefs:
	typedef Key key_type;
	typedef T mapped_type;
	typedef pair<const Key, T> value_type;
private:
	typedef rb_tree<key_type, value_type,
		select1st<value_type>, key_compare, Alloc> rep_type;
	rep_type t; // red-black tree representing map
};
// stl_tree.h
struct __rb_tree_node_base
{
	typedef __rb_tree_color_type color_type;
	typedef __rb_tree_node_base* base_ptr;
	color_type color;
	base_ptr parent;
	base_ptr left;
	base_ptr right;
};
// stl_tree.h
template <class Key, class Value, class KeyOfValue, class Compare, class Alloc
	= alloc>
class rb_tree {
protected:
	typedef void* void_pointer;
	typedef __rb_tree_node_base* base_ptr;
	typedef __rb_tree_node<Value> rb_tree_node;
	typedef rb_tree_node* link_type;
	typedef Key key_type;
	typedef Value value_type;
public:
	// insert⽤的是第⼆个模板参数左形参
	pair<iterator, bool> insert_unique(const value_type& x);
	// erase和find⽤第⼀个模板参数做形参
	size_type erase(const key_type& x);
	iterator find(const key_type& x);
protected:
	size_type node_count; // keeps track of size of tree
	link_type header;
};
template <class Value>
struct __rb_tree_node : public __rb_tree_node_base
{
	typedef __rb_tree_node<Value>* link_type;
	Value value_field;
};

1.2 对比set和map的源码：泛型编程的应用

虽然底层都是用红黑树实现的，这里我们看源码，第一个模板参数都是Key，区别就在于第二个模板参数，value对于set是key，对map不是------对于map是一个pair<const Key, T>

map和set不是同一棵红黑树实现的，这里其实是用红黑树类模板实现的。

我们通过上面的图片框架可以看到源码中的rb_tree实现了一个非常巧妙的泛型思想，rb_tree不管是实现Key的搜索场景还是Key / value的搜索场景不是直接写死的，而是由第二个模板参数value来决定_rb_tree_node中存储的数据类型。

set实例化rb_tree的时候第二个模板参数给的是Key，map实例化rb_tree时第二个模板参数给的是pair<const key , T>，这样一个红黑树既可以实现Key搜索场景的set，也可以实现Key / value搜索场景的map。

注意，源码里面模板参数是用了T来代表value，而内部写的value_type不是我们日常Key / value场景中说的value，源码中的value_type反而是红黑树节点中，存储的是真实的数据类型。

rb_tree第二个模板参数value已经控制了红黑树节点存储的数据类型，为什么还要穿第一个模板参数Key呢？尤其是set，两个模板参数是一样的，这是很多uu会有的有个疑问。要注意的是对于map和set，find / erase时候的模板参数都是Key，所以第一个模板参数是传给find / erase等函数来做形参的类型的。对于set而言，两个参数是一样的；但是对于map而言就不一样了，map容器Insert的是pair对象，但是find / erase的是Key对象。

2. 模拟实现map和set

2.1 实现出复用红黑树的框架，并支持insert

参考前面的源码框架，map和set确实是复用之前我们实现的红黑树。

这里相比源码调整一下，key参数就用K，value参数就用V，红黑树中的数据类型，我们使用T。

其次因为RBTree实现了泛型不知道T参数导致是K，还是pair<K，V>，那么insert内部进行插入逻辑比较时，就没办法进行比较，因为pair的默认支持的是key和value一起参与比较，我们需要时的任何时候只比较key，所以我们在map和set层分别实现一个MapKeyOfT和SetKeyOfT的仿函数传给

RBTree的KeyOfT，然后RBTree中通过KeyOfT仿函数取出T类型对象中的key，再进行比较，具体

细节参考如下代码实现。

Map.h

cpp 复制代码

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

Set.h

cpp 复制代码

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

RBTree.h

实现步骤：

实现红黑树

封装map和set框架，解决keyOfT

iterator

const_iterator

key不支持修改的问题

operator[]

cpp 复制代码

enum Colour
{
	RED,
	BLACK
};
template<class T>
struct RBTreeNode
{
	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)
	{}
};

template<class K, class T, class KeyOfT>
class RBTree
{
private:
	typedef RBTreeNode<T> Node;
	Node* _root = nullptr;
public:
	bool Insert(const T& data)
	{
		if (_root == nullptr)
		{
			_root = new Node(data);
			_root->_col = BLACK;
			return true;
        }
        		KeyOfT kot;
		Node* parent = nullptr;
		Node* cur = _root;
		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);
		Node* newnode = cur;
		// 新增结点。颜⾊给红色
		cur->_col = RED;
		if (kot(parent->_data) < kot(data))
		{
			parent->_right = cur;
		}
		else
		{
			parent->_left = cur;
		}
		cur->_parent = parent;
		//...
		return true;
	}
}

2.2 支持iterator的实现

iterator核心源代码

支持完整的迭代器还有很多细节需要修改，具体参考下面题的代码。

2.3 map支持[]

map要⽀持[]主要需要修改insert返回值支持，修改RBtree中的insert返回值为
pair<Iterator, bool> Insert(const T& data)
有了insert支持[]实现就很简单了，具体参考下面代码实现

2.4 bit::map和bit::set代码实现

Set.h

cpp 复制代码

#pragma once

#include"RBTree.h"

namespace bit
{
	template<class K>
	class set
	{
		struct SetKeyOfT
		{
			const K& operator()(const K& key)
			{
				return key;
			}
		};
	public:
		typedef typename RBTree<K, const K, SetKeyOfT>::Iterator iterator;
		typedef typename RBTree<K, const K, SetKeyOfT>::ConstIterator const_iterator;

		iterator begin()
		{
			return _t.Begin();
		}

		iterator end()
		{
			return _t.End();
		}

		const_iterator begin() const
		{
			return _t.Begin();
		}

		const_iterator end() const
		{
			return _t.End();
		}

		pair<iterator, bool> insert(const K& key)
		{
			return _t.Insert(key);
		}

		iterator find(const K& key)
		{
			return _t.Find(key);
		}
		
	private:
		RBTree<K, const K, SetKeyOfT> _t;
	};
}

Map.h

cpp 复制代码

#pragma once

#include"RBTree.h"
namespace bit
{
	template<class K, class V>
	class map
	{
		struct MapKeyOfT
		{
			const K& operator() (const pair<K, V>& kv)
			{
				return kv.first;
			}
		};
	public:
		typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::Iterator iterator;
		typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::ConstIterator const_iterator;

		iterator begin()
		{
			return _t.Begin();
		}

		iterator end()
		{
			return _t.End();
		}

		const_iterator begin() const
		{
			return _t.Begin();
		}

		const_iterator end() const
		{
			return _t.End();
		}

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

		iterator find(const K& key)
		{
			return _t.Find(key);
		}

		V& operator[](const K& key)
		{
			//pair<iterator, bool> ret = _t.Insert({key, V() });
			auto [it, flag] = _t.Insert({ key, V() });
			return it->second;
		}

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

RBTree.h

cpp 复制代码

#pragma once

#pragma once
#include<assert.h>

enum Colour
{
	RED,
	BLACK
};

// red black tree
template<class T>
struct RBTreeNode
{
	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)
		, _col(RED)
	{}
};

// RBTree<K, pair<K, V>> _t;-> // map
// RBTree<K, K> _t;->		   // set

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(node)
	{}

	//解引用迭代器，返回当前节点存储的数据的引用
	Ref operator*()
	{
		return _node->_data;
	}

	//返回当前数据的地址（指针），用于访问对象的成员
	Ptr operator->()
	{
		return &_node->_data;
	}

	// 返回下一个对象的引用
	Self& operator++()
	{
		// 如果存在右子树，下一个找到右子树的最左节点
		if (_node->_right)
		{
			Node* minRight = _node->_right;
			while (minRight->_left)
			{
				minRight = minRight->_left;
			}
			_node = minRight;
		}

		// 如果右子树不存在，想上找祖先结点，直到找到孩子为父亲的左的祖先节点
		else
		{
			Node* cur = _node;
			Node* parent = cur->_parent;
			while (parent && parent->_right == cur)
			{
				cur = parent;
				parent = cur->_parent;
			}
			_node = parent;
		}
		return *this;
	}

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

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

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

	~RBTree()
	{
		Destroy(_root);
		_root = nullptr;
	}

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

	Iterator End()
	{
		return Iterator(nullptr);
	}

	ConstIterator Begin() const
	{
		Node* minLeft = _root;
		while (minLeft && minLeft->_left)
		{
			minLeft = minLeft->_left;
		}
		return ConstIterator(minLeft);
	}

	ConstIterator End() const
	{
		return ConstIterator(nullptr);
	}

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

		KeyOfT kot;
		Node* parent = nullptr;
		Node* cur = _root;
		while (cur)
		{
			if (kot(cur->_data) < kot(data))
			{
				parent = cur;
				cur = cur->_right;
			}
			//else if (kot(cur->_data) > kot(data))
			else if (kot(data) < kot(cur->_data))
			{
				parent = cur;
				cur = cur->_left;
			}
			else
			{
				return {Iterator(cur), false};
			}
		}

		// 新增红色
		cur = new Node(data);
		Node* newnode = cur;  // newnode为新插入的结点
		cur->_col = RED;
		if (kot(parent->_data) < kot(data))
		{
			parent->_right = cur;
		}
		else
		{
			parent->_left = cur;
		}
		cur->_parent = parent;

		// 8:10
		while (parent && parent->_col == RED)
		{
			Node* grandfather = parent->_parent;
			if (grandfather->_left == parent)
			{
				//    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 // 叔叔不存在或者叔叔存在且为黑
				{
					//    g
					//  p   u
					//c
					// 单旋+变色
					if (cur == parent->_left)
					{
						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
					{
						//     g
						//   u   p
						// c
						RotateR(parent);
						RotateL(grandfather);
						cur->_col = BLACK;
						grandfather->_col = RED;
					}
					break;
				}
			}
		}

		_root->_col = BLACK;

		return {Iterator(newnode), true};
	}

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

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

		Node* parentParent = parent->_parent;

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

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

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

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

		return End();
	}


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

	int Size()
	{
		return _Size(_root);
	}
private:
	int _Size(Node* root)
	{
		if (root == nullptr)
			return 0;

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

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


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

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

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

		Destroy(root->_left);
		Destroy(root->_right);
		delete root;
	}

private:
	Node* _root = nullptr;
};

Test.cpp

cpp 复制代码

#include"Map.h"
#include"Set.h"

template<class T>
void func(const bit::set<T>& s)
{
	typename bit::set<T>::const_iterator it = s.begin();
	while (it != s.end())
	{
		//*it = 1;
		cout << *it << " ";
		++it;
	}
	cout << endl;
}

void test_set()
{
	bit::set<int> s;
	s.insert(1);
	s.insert(2);
	s.insert(1);
	s.insert(5);
	s.insert(0);
	s.insert(10);
	s.insert(8);

	bit::set<int>::iterator it = s.begin();
	// *it += 10;
	while (it != s.end())
	{
		cout << *it << " ";
		++it;
	}
	cout << endl;

	func(s);
}

void test_map()
{
	bit::map<string, string> dict;
	dict.insert({ "sort", "排序" });
	dict.insert({ "left", "左边" });
	dict.insert({ "right", "右边" });

	dict["string"] = "字符串"; // 插入+修改
	dict["left"] = "左边xxx";  // 修改

	auto it = dict.begin();
	while (it != dict.end())
	{
		// it->first += 'x'; // 不能修改
		it->second += 'x';

		cout << it->first << ":" << it->second << endl;
		++it;
	}
	cout << endl;

	for (auto& [k, v] : dict)
	{
		cout << k << ":" << v << endl;
	}
	cout << endl;

	string arr[] = { "苹果", "西瓜", "苹果", "西瓜", "苹果", "苹果", "西瓜", "苹果", "香蕉", "苹果", "香蕉" };
	bit::map<string, int> countMap;
	for (auto& e : arr)
	{
		/*auto it = countMap.find(e);
		if (it != countMap.end())
		{
			it->second++;
		}
		else
		{
			countMap.insert({ e, 1 });
		}*/
		countMap[e]++;
	}

	for (auto& [k, v] : countMap)
	{
		cout << k << ":" << v << endl;
	}
	cout << endl;
}

int main()
{
	test_set();
	test_map();

	return 0;
}

运行结果：

【注意：typename的用法】

typename是 C++ 中一个用于模板的关键词 ，主要解决依赖类型（dependent types）的歧义问题 。它的核心作用是：告诉编译器"这是一个类型名，不是变量"。

只要在模板里看到 T::XXX，且 XXX 是个类型，就加 typename