这篇文章主要是讲解 map 和 set 的模拟实现。
1 map 与 set 源码框架分析
我们先来看一下 stl 源码中 map 和 set 是如何实现的（下面为 stl 源码中的部分核心代码）：
cpp
//stl_set.h
template<typename _Key, typename _Compare = std::less<_Key>,
typename _Alloc = std::allocator<_Key> >
class set
{
public:
typedef _Key key_type;
typedef _Key value_type;
typedef _Compare key_compare;
typedef _Compare value_compare;
typedef _Alloc allocator_type;
private:
typedef _Rb_tree<key_type, value_type, _Identity<value_type>, key_compare, _Key_alloc_type> _Rep_type;
_Rep_type _M_t;
};
//stl_map.h
template <typename _Key, typename _Tp, typename _Compare = std::less<_Key>,
typename _Alloc = std::allocator<std::pair<const _Key, _Tp> > >
class map
{
public:
typedef _Key key_type;
typedef _Tp mapped_type;
typedef std::pair<const _Key, _Tp> value_type;
typedef _Compare key_compare;
typedef _Alloc allocator_type;
private:
typedef _Rb_tree<key_type, value_type, _Select1st<value_type>,
key_compare, _Pair_alloc_type> _Rep_type;
_Rep_type _M_t;
};
//stl_tree.h
enum _Rb_tree_color { _S_red = false, _S_black = true };
struct _Rb_tree_node_base
{
typedef _Rb_tree_node_base* _Base_ptr;
typedef const _Rb_tree_node_base* _Const_Base_ptr;
_Rb_tree_color _M_color;
_Base_ptr _M_parent;
_Base_ptr _M_left;
_Base_ptr _M_right;
};
template<typename _Key, typename _Val, typename _KeyOfValue,
typename _Compare, typename _Alloc = allocator<_Val> >
class _Rb_tree
{
protected:
typedef _Rb_tree_node_base* _Base_ptr;
typedef const _Rb_tree_node_base* _Const_Base_ptr;
public:
typedef _Key key_type;
typedef _Val value_type;
typedef value_type* pointer;
typedef const value_type* const_pointer;
typedef value_type& reference;
typedef const value_type& const_reference;
typedef _Rb_tree_node<_Val>* _Link_type;
typedef const _Rb_tree_node<_Val>* _Const_Link_type;
typedef size_t size_type;
typedef ptrdiff_t difference_type;
typedef _Alloc allocator_type;
};可以看到在源码中，map 和 set 是复用的 stl_tree.h 中的 _Rb_tree 来实现的，也就是对 _Rb_tree 类的进一步封装，而且最重要的是复用的一棵 _Rb_tree 来实现的。具体来说，库中的 map 和 set 实现了进一步的泛型，使得一颗红黑树即可完成 map 和 set 的封装。那么具体是怎么实现的呢？
按照我们之前的实现思路，RBtree 中一共有两个模板参数，一个是 K，一个是 V，K 是 key 的模板参数，V 是 value 的模板参数，但是在库中的 _Rb_tree 虽然也有两个模板参数 _Key 与 _Val，但是如果我们仔细看 map 和 set 对于 _Rb_tree 的封装，可以发现，_Key 是键值的模板参数，但是 _Val 却是每个节点中存储的数据类型，也就是对于 set 来说 _Val 中依旧是 Key，但是对于 map 来说，_Val 却是 pair<K, V>。所以将库中的 map 和 set 简化来写就是这样的：
cpp
//由于 set 中放的是 key, map 中放的是 pair<K, V>, 所以这里写为 T
template<class T>
struct RBNode
{
T _data;
RBNode<T>* _left;
RBNode<T>* _right;
RBNode<T>* _parent;
Color _col;
};
//第二个类型对于 set 是 K, 对于 map 是 pair<k, V>, 所以这里写为 T
template<class K, class T>
class RBTree
{
typedef RBNode<T> Node;
Node* _root;
};
template<class K>
class set
{
RBTree<K, K> _t;
};
template<class V>
class map
{
RBTree<K, pair<K, V>> _t;
};库中这样设计最大的好处就是进一步提升了泛型编程，实现了代码的进一步复用。库中将 V 设计为 map 或者 set 中实际存储的数据类型，map 和 set 就都可以通过封装这一棵红黑树实现，我们就不用实现两棵逻辑相同的红黑树了，在模板之上进一步提升了泛型。
虽然库中实现了进一步的泛型，但是也带来了一个麻烦的地方，就是必须添加一个模板参数，也就是红黑树中的第一个 K 模板参数。初看这个 K 模板参数，对于 set 来说，简直就是一个多余的参数，因为这个 K 的存在，所有必须传递两个相同的模板参数，也就是 RBTree<K, K>，但是传递这个参数是不得已而为之的存在。因为 map 和 set 都需要通过封装 RBTree 来实现，所以第二个模板参数 V 就不确定是 K 还是 pair<K, V>，但是在 find 或者 erase 函数中，我们又需要用到 key，那如果只有 V 这个模板参数，那么我们是用 V 直接比较（set）还是用 V.first 比较（map） 呢？所以我们必须通过 K 这一模板参数来实现 find 或者 erase 函数。
2 封装 RBTree 来实现 map 与 set
我们就按照源码中的框架来修改我们之前的 RBTree，然后对其封装来实现 map 和 set。
修改 RBTree 与封装 RBTree
按照源码中的思路，我们需要用到两个模板参数，这里我们将两个模板参数设计为 K 和 T，K 代表 key，T 对于 set 来说是 key，对于 map 来说是 pair<K, V>：
cpp
enum Color
{
BLACK,
RED
};
//节点的类型也要修改为 T
template<class T>
struct RBNode
{
T _data;
RBNode<T>* _left;
RBNode<T>* _right;
RBNode<T>* _parent;
Color _col;
RBNode(const T& data)
:_data(data)
,_left(nullptr)
, _right(nullptr)
,_parent(nullptr)
,_col(BLACK)
{}
};
//K -- key，T -- key/pair<K, V>
template<class K, class T>
class RBTree
{
typedef RBNode<T> Node;
Node* _root = nullptr;
};有了 RBTree 我们就可以封装这一棵红黑树来实现 map 和 set，在实现时，为了防止与库中的 map 和 set 冲突，所以我们选择使用命名空间封装起来：
cpp
//myset.hpp
#pragma once
#include "RBTree.hpp"
namespace LTL
{
template<class K>
class set
{
public:
private:
RBTree<K, K> _t;
};
}
//mymap.hpp
#pragma once
#include "RBTree.hpp"
namespace LTL
{
template<class K, class V>
class map
{
public:
private:
RBTree<K, pair<K, V>> _t;
};
}KeyOfT 模板参数
上面我们已经将框架修改为源码中的框架了，但是我们在实现修改后的 insert 函数时，发现遇到困难了，比较 data 不知道是直接比较还是 data.first 比较了：
cpp
bool Insert(const T& data)
{
if (_root == nullptr)
{
_root = new Node(data);
return true;
}
Node* cur = _root;
Node* parent = nullptr;
while (cur)
{
//在这里进行键值比较时, data 可能是 K, 也有可能是 pair<K, V>, 所以不知道如何比较
if (cur->_data.first > data.first)
{
parent = cur;
cur = cur->_left;
}
else if (cur->_data.first < data.first)
{
parent = cur;
cur = cur->_right;
}
else
return false;
}
cur = new Node(data);
if (parent->_data.first > data.first)
parent->_left = cur;
else
parent->_right = cur;
//不要忘记链接 _parent
cur->_parent = parent;
//新插入节点的颜色为红色
cur->_col = RED;
//父亲不为空，并且 parent 颜色为红色时才需要更新颜色
while (parent && parent->_col == RED)
{
//parent 为红色，grandparent 一定存在
Node* grandparent = parent->_parent;
if (parent == grandparent->_left)
{
Node* uncle = grandparent->_right;
//如果 uncle 存在且为红色
if (uncle && uncle->_col == RED)
{
//将 uncle 与 parent 变黑，grandparent 变红
uncle->_col = parent->_col = BLACK;
grandparent->_col = RED;
//继续向上更新
cur = grandparent;
parent = cur->_parent;
}
else
{
//uncle 不存在或者 uncle 存在且为黑
if (cur == parent->_left)
{
//不管哪种情况，只要在左边，就进行右单旋
RotateR(grandparent);
//将 parent 更新为黑色，grandparent 更新为红色
parent->_col = BLACK;
grandparent->_col = RED;
}
else
{
//进行左右双旋
RotateLR(grandparent);
//将 cur 更新为黑色，grandparent 更新为红色
cur->_col = BLACK;
grandparent->_col = RED;
}
//旋转之后，不需要更新了
break;
}
}
else
{
Node* uncle = grandparent->_left;
//如果 uncle 存在且为红色
if (uncle && uncle->_col == RED)
{
//将 uncle 与 parent 变黑，grandparent 变红
uncle->_col = parent->_col = BLACK;
grandparent->_col = RED;
//继续向上更新
cur = grandparent;
parent = cur->_parent;
}
else
{
//uncle 不存在或者 uncle 存在且为黑
if (cur == parent->_right)
{
//不管哪种情况，只要在右边，就进行左单旋
RotateL(grandparent);
//将 parent 更新为黑色，grandparent 更新为红色
parent->_col = BLACK;
grandparent->_col = RED;
}
else
{
//进行右左双旋
RotateRL(grandparent);
//将 cur 更新为黑色，grandparent 更新为红色
cur->_col = BLACK;
grandparent->_col = RED;
}
//旋转之后，不需要继续更新了
break;
}
}
}
//根节点可能被更新为红色，要重新变为黑色
_root->_col = BLACK;
return true;
}产生这种原因的根本原因就是底层的 RBTree 并不知道 data 是什么，但是上层的 set 和 map 是知道的呀，所以我们可以在 RBTree 中加入一个 KeyOfT 模板参数，用来取出 data 中的 key，这样上层的 set 和 map 传入对应的仿函数，这样就可以实现 data 的比较逻辑了：
cpp
//RBTree.hpp
//K -- key，T -- key/pair<K, V>
//添加 KeyOfT 模板参数, 取出 data 中的 key 值
template<class K, class T, class KeyOfT>
class RBTree
{
typedef RBNode<T> Node;
public:
bool Insert(const T& data)
{
if (_root == nullptr)
{
_root = new Node(data);
return true;
}
Node* cur = _root;
Node* parent = nullptr;
KeyOfT kot;
while (cur)
{
if (kot(cur->_data) > kot(data))
{
parent = cur;
cur = cur->_left;
}
else if (kot(cur->_data) < kot(data))
{
parent = cur;
cur = cur->_right;
}
else
return false;
}
cur = new Node(data);
if (kot(parent->_data) > kot(data))
parent->_left = cur;
else
parent->_right = cur;
//不要忘记链接 _parent
cur->_parent = parent;
//新插入节点的颜色为红色
cur->_col = RED;
//父亲不为空，并且 parent 颜色为红色时才需要更新颜色
while (parent && parent->_col == RED)
{
//parent 为红色，grandparent 一定存在
Node* grandparent = parent->_parent;
if (parent == grandparent->_left)
{
Node* uncle = grandparent->_right;
//如果 uncle 存在且为红色
if (uncle && uncle->_col == RED)
{
//将 uncle 与 parent 变黑，grandparent 变红
uncle->_col = parent->_col = BLACK;
grandparent->_col = RED;
//继续向上更新
cur = grandparent;
parent = cur->_parent;
}
else
{
//uncle 不存在或者 uncle 存在且为黑
if (cur == parent->_left)
{
//不管哪种情况，只要在左边，就进行右单旋
RotateR(grandparent);
//将 parent 更新为黑色，grandparent 更新为红色
parent->_col = BLACK;
grandparent->_col = RED;
}
else
{
//进行左右双旋
RotateLR(grandparent);
//将 cur 更新为黑色，grandparent 更新为红色
cur->_col = BLACK;
grandparent->_col = RED;
}
//旋转之后，不需要更新了
break;
}
}
else
{
Node* uncle = grandparent->_left;
//如果 uncle 存在且为红色
if (uncle && uncle->_col == RED)
{
//将 uncle 与 parent 变黑，grandparent 变红
uncle->_col = parent->_col = BLACK;
grandparent->_col = RED;
//继续向上更新
cur = grandparent;
parent = cur->_parent;
}
else
{
//uncle 不存在或者 uncle 存在且为黑
if (cur == parent->_right)
{
//不管哪种情况，只要在右边，就进行左单旋
RotateL(grandparent);
//将 parent 更新为黑色，grandparent 更新为红色
parent->_col = BLACK;
grandparent->_col = RED;
}
else
{
//进行右左双旋
RotateRL(grandparent);
//将 cur 更新为黑色，grandparent 更新为红色
cur->_col = BLACK;
grandparent->_col = RED;
}
//旋转之后，不需要继续更新了
break;
}
}
}
//根节点可能被更新为红色，要重新变为黑色
_root->_col = BLACK;
return true;
}
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;
}
return nullptr;
}
private:
Node* _root = nullptr;
};
//myset.hpp
#pragma once
#include "RBTree.hpp"
namespace LTL
{
template<class K>
class set
{
public:
struct SetKeyOfT
{
K& operator()(const K& key)
{
return key;
}
};
private:
RBTree<K, K, SetKeyOfT> _t;
};
}
//mymap.hpp
#pragma once
#include "RBTree.hpp"
namespace LTL
{
template<class K, class V>
class map
{
public:
struct MapKeyOfT
{
K& operator()(const pair<K, V>& kv)
{
return kv.first;
}
};
private:
RBTree<K, pair<K, V>, MapKeyOfT> _t;
};
}实现 iterator 与 const_iterator
在红黑树中，比较麻烦的就是实现迭代器了。在 map 和 set 的使用中，我们知道其迭代器遍历是走的整棵树的中序遍历，在数据结构初阶中，我们是通过递归实现的二叉树的中序遍历，但是这里我们需要通过迭代，也就是循环的方式来实现红黑树的中序遍历（也就是迭代器中的 operator++ 与 operator--）。
1）前向中序遍历（operator++）
首先我们先来看一下一棵二叉树是如何向前进行中序遍历的。首先，我们都知道中序遍历的顺序是左根右，也就是左子树 -> 根节点 -> 右子树，所以对于一棵红黑树来说，中序遍历的起始节点就是整棵树的最左节点。找到了最开始的节点，我们就以最开始的节点为例，我们就来模仿一下中序遍历的过程：
其中 it 采用中序遍历每个节点。我们就把每个节点想成以当前节点为子树的根节点，那么遵循左根右的遍历原则，如果当前节点的右子树一旦不为空，那么中序遍历的下一个节点必然位于该节点的右子树中，那么是右子树的哪个节点呢？我们可以看图中的 30 节点：
当 it 遍历到 30 之后，其遍历的下一个节点就是 35，就是其右子树的最左节点，也就是右子树中中序遍历的第一个节点。所以对于 20、30、45 这些右子树不为空的节点，其下一个节点无一例外都是走到右子树的最左节点。
那么如果当前节点的右子树为空了呢？比如 15、25、50 这些节点，我们按照中序遍历的思想，一旦一个根节点的右子树为空了，那就说明以当前节点为根节点的整棵子树就已经遍历完了，如果当前节点还是其父亲节点的右孩子节点，那么以及父亲节点为根节点的整棵子树也遍历完了，如果其父亲节点还是其爷爷节点的右孩子节点，那就说明以及爷爷节点的整棵子树也遍历完了，于是就一步一步向上寻找，那么寻找到哪个祖先节点呢？我们可以看 15 这个节点，其右子树为空，但是由于该节点是其父亲节点的左孩子节点，所以说明 20 的左孩子遍历完了，根据左根右的遍历顺序，接下来就该遍历当前子树的根节点，也就是 15 的父亲节点 20 了。所以就是寻找到孩子为其父亲节点的左孩子的那个父亲节点即可，因为下一个节点就是这个父亲节点。
上面的过程可能比较抽象，我们举一个例子。就比如上面图中的 25 节点：
25 节点遍历完了，由于 25 是 20 的右孩子节点，所以说明 20 这棵子树遍历结束了；但是 20 是 30 的左子树，说明 30 的左子树遍历完了，下一个节点正好就是 30 这棵子树的根节点，也就是 30 节点，所以 it 下一个节点就会到 30 这个节点。
总结一下，其实向前的中序遍历对于每个节点就分为两种情况：
（1） 右子树不为空：去到的下一个节点为右子树的最左节点
（2） 右子树为空：去到的下一个节点为当前节点是其父亲节点的左孩子节点的父亲节点（是祖先节点中的其中一个）
2）后向中序遍历（operator--）
后向中序遍历就是前向中序遍历的反过程，也就是右根左的过程，那么根据上面前向中序遍历的过程，后向中序遍历其实就是分为两种情况：
（1） 左子树不为空：去到的下一个节点为左子树的最右节点
（2） 左子树为空：去到的下一个节点为当前节点是其父亲节点的右孩子节点的父亲节点（是祖先节点中的其中一个）
3）实现迭代器
有了使用迭代的方式实现中序遍历的算法思想，我们就可以来实现 map 和 set 的迭代器了。首先 map 和 set 的迭代器也必须和 list 中的一样，必须创建一个 iterator 类，在其中重载对应的操作符才能实现，要不然 RBTree 中的节点为指针，--、++ 迭代器操作就是对指针--、++，很显然达不到我们想要的效果。
与链表的迭代器类似，迭代器中的成员变量就是 RBNode<T>* _node，另外我们在这里添加一个 RBNode<T>* _root 成员变量来指向红黑树中的根节点，这样做主要是为了方便我们实现 operator--，因为在迭代器中是可以对 end() 做 -- 的，对 end() 做 --，我们需要找到树的根节点，如果我们不传入 _root，我们无法找到根节点。
cpp
//RBTree.hpp
//...
//与 list 相同，Ptr、Ref 模板参数主要是方便实现 iterator 或者 const_iterator 的
template<class T, class Ptr, class Ref>
struct _Rb_Iterator
{
typedef RBNode<T> Node;
typedef _Rb_Iterator<T, Ptr, Ref> Self;
Node* _node;
//这里有一个 _root 主要是方便 end() 做--
Node* _root;
_Rb_Iterator(Node* node, Node* root)
:_node(node)
,_root(root)
{}
Ref operator*()
{
return _node->_data;
}
Ptr operator->()
{
return &_node;
}
bool operator!=(const Self& it)
{
return _node != it._node;
}
//前置++
Self& operator++()
{
//1. 如果当前节点的右孩子不为空
if (_node->_right)
{
//去到右孩子的最左节点
_node = _node->_right;
while (_node->_left)
{
_node = _node->_left;
}
}
//2. 右孩子为空
else
{
//去到当前节点为父亲节点左孩子的父亲节点
Node* cur = _node;
Node* parent = cur->_parent;
//parent 可能为空，比如图中的 50 节点
while (parent && cur == parent->_right)
{
cur = parent;
parent = parent->_parent;
}
_node = parent;
}
return *this;
}
//后置++
Self operator++(int)
{
Self tmp = *this;
++(*this);
return tmp;
}
//前置--
Self& operator--()
{
//operator--中可以对 end() 进行--
if (_node == nullptr)
{
_node = _root;
return *this;
}
//与 operator++ 的逻辑相反
//1. 左孩子不为空
if (_node->_left)
{
//去到左孩子的最右节点
_node = _node->_left;
while (_node->_right)
{
_node = _node->_right;
}
}
//2. 左孩子为空
else
{
//去到当前节点为父亲节点右孩子的父亲节点
Node* cur = _node;
Node* parent = cur->_parent;
//parent 可能为空，比如图中的 50 节点
while (parent && cur == parent->_left)
{
cur = parent;
parent = parent->_parent;
}
_node = parent;
}
return *this;
}
//后置--
Self operator--(int)
{
Self tmp = *this;
--(*this);
return tmp;
}
};
//K -- key，T -- key/pair<K, V>
template<class K, class T, class KeyOfT>
class RBTree
{
typedef RBNode<T> Node;
//...
public:
typedef _Rb_Iterator<T, T*, T&> Iterator;
typedef _Rb_Iterator<T, const T*, const T&> Const_Iterator;
Iterator Begin()
{
//找到最左节点
Node* cur = _root;
while (cur->_left)
{
cur = cur->_left;
}
return Iterator(cur, _root);
}
Iterator End()
{
return Iterator(nullptr, _root);
}
Const_Iterator Begin() const
{
//找到最左节点
Node* cur = _root;
while (cur->_left)
{
cur = cur->_left;
}
return Const_Iterator(cur, _root);
}
Const_Iterator End() const
{
return Const_Iterator(nullptr, _root);
}
//...
private:
Node* _root = nullptr;
};
//myset.hpp
#pragma once
#include "RBTree.hpp"
namespace LTL
{
template<class K>
class set
{
public:
struct SetKeyOfT
{
K& operator()(const K& key)
{
return key;
}
};
//实现迭代器
//这里需要在前面加 typename 关键字表明 Iterator 是个类型
typedef typename RBTree<K, K, SetKeyOfT>::Iterator iterator;
typedef typename RBTree<K, K, SetKeyOfT>::Const_Iterator 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();
}
private:
RBTree<K, K, SetKeyOfT> _t;
};
}
//mymap.hpp
#pragma once
#include "RBTree.hpp"
namespace LTL
{
template<class K, class V>
class map
{
public:
struct MapKeyOfT
{
K& operator()(const pair<K, V>& kv)
{
return kv.first;
}
};
//这里需要在前面加 typename 关键字表明 Iterator 是个类型
typedef typename RBTree<K, pair<K, V>, MapKeyOfT>::Iterator iterator;
typedef typename RBTree<K, pair<K, V>, MapKeyOfT>::Const_Iterator 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();
}
private:
RBTree<K, pair<K, V>, MapKeyOfT> _t;
};
}Key 不可修改
我们虽然已经实现了大部分了，但是有一个致命的问题，就是 set 与 map 的 Key 是可以被修改的，但是现实中 Key 要被用作键值，Key 要用作比较逻辑，一旦 Key 被修改，那么整棵红黑树的结构就变了，所以 Key 是不能修改的。
实现 Key 不能修改的逻辑也很简单，只要在封装的时候，对 K 模板参数前加入 const 关键字即可。对于 set 来说，底层的红黑树中的 V 就是 const K，data 的类型就是 const K，也就无法修改了；对于 map 来说，Key 不能修改，但是 Value 是可以修改的，所以我们只需要将 RBTree 中 V 模板参数传入 pair<const K, V>，这样 data 的类型就是 pair<const K, V>，这样 Key 就无法修改了，只有 Value 可以修改。
cpp
//myset.hpp
#pragma once
#include "RBTree.hpp"
namespace LTL
{
template<class K>
class set
{
public:
struct SetKeyOfT
{
//这里的 K 不要忘记修改为 const K& 哦
const K& operator()(const K& key)
{
return key;
}
};
//实现迭代器
//这里需要在前面加 typename 关键字表明 Iterator 是个类型
typedef typename RBTree<K, const K, SetKeyOfT>::Iterator iterator;
typedef typename RBTree<K, const K, SetKeyOfT>::Const_Iterator const_iterator;
//...
private:
//将 K 修改为 const 即可
//这里修改了，别忘记上面 typedef 部分也要修改
RBTree<K, const K, SetKeyOfT> _t;
};
}
//mymap.hpp
#pragma once
#include "RBTree.hpp"
namespace LTL
{
template<class K, class V>
class map
{
public:
struct MapKeyOfT
{
//这里的 K 不要忘记修改为 const K& 哦
const K& operator()(const pair<K, V>& kv)
{
return kv.first;
}
};
//这里需要在前面加 typename 关键字表明 Iterator 是个类型
typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::Iterator iterator;
typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::Const_Iterator const_iterator;
//...
private:
//只将 key 修改为 const 即可
RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};
}实现 map 中的 operator[] 功能
在使用阶段我们使用过这个 operator[] 接口，这个方法的功能是如果键值不存在，那就可以插入键值对，如果插入的键值存在，那么我们就可以修改 value，也可以查找对应的 value 值，那么这个接口是如何做到的呢？
这个接口其实是依托于 insert 接口实现的，insert 的返回值其实为 pair<iterator, bool>，所以一旦 key 不存在，那么 insert 就会返回插入键值对之后的迭代器，在上层的 operator[] 也就完成了插入；一旦存在，那么 insert 就会返回已经存在的键值对的迭代器，所以就可以查找到对应的键值对了，之所以能够修改，就是因为 operator[] 返回的是 V&，这样就可以修改了。
cpp
//RBTree.hpp
//K -- key，T -- key/pair<K, V>
template<class K, class T, class KeyOfT>
class RBTree
{
typedef RBNode<T> Node;
//...
public:
//...
//修改 Insert 返回值为 pair<Iterator, bool>
pair<Iterator, bool> Insert(const T& data)
{
if (_root == nullptr)
{
_root = new Node(data);
return { Iterator(_root, _root), true };
}
Node* cur = _root;
Node* parent = nullptr;
KeyOfT kot;
while (cur)
{
if (kot(cur->_data) > kot(data))
{
parent = cur;
cur = cur->_left;
}
else if (kot(cur->_data) < kot(data))
{
parent = cur;
cur = cur->_right;
}
else
return { Iterator(cur, _root), false };
}
cur = new Node(data);
//保存一下新插入的节点
Node* newnode = cur;
if (kot(parent->_data) > kot(data))
parent->_left = cur;
else
parent->_right = cur;
//不要忘记链接 _parent
cur->_parent = parent;
//新插入节点的颜色为红色
cur->_col = RED;
//父亲不为空，并且 parent 颜色为红色时才需要更新颜色
while (parent && parent->_col == RED)
{
//parent 为红色，grandparent 一定存在
Node* grandparent = parent->_parent;
if (parent == grandparent->_left)
{
Node* uncle = grandparent->_right;
//如果 uncle 存在且为红色
if (uncle && uncle->_col == RED)
{
//将 uncle 与 parent 变黑，grandparent 变红
uncle->_col = parent->_col = BLACK;
grandparent->_col = RED;
//继续向上更新
cur = grandparent;
parent = cur->_parent;
}
else
{
//uncle 不存在或者 uncle 存在且为黑
if (cur == parent->_left)
{
//不管哪种情况，只要在左边，就进行右单旋
RotateR(grandparent);
//将 parent 更新为黑色，grandparent 更新为红色
parent->_col = BLACK;
grandparent->_col = RED;
}
else
{
//进行左右双旋
RotateLR(grandparent);
//将 cur 更新为黑色，grandparent 更新为红色
cur->_col = BLACK;
grandparent->_col = RED;
}
//旋转之后，不需要更新了
break;
}
}
else
{
Node* uncle = grandparent->_left;
//如果 uncle 存在且为红色
if (uncle && uncle->_col == RED)
{
//将 uncle 与 parent 变黑，grandparent 变红
uncle->_col = parent->_col = BLACK;
grandparent->_col = RED;
//继续向上更新
cur = grandparent;
parent = cur->_parent;
}
else
{
//uncle 不存在或者 uncle 存在且为黑
if (cur == parent->_right)
{
//不管哪种情况，只要在右边，就进行左单旋
RotateL(grandparent);
//将 parent 更新为黑色，grandparent 更新为红色
parent->_col = BLACK;
grandparent->_col = RED;
}
else
{
//进行右左双旋
RotateRL(grandparent);
//将 cur 更新为黑色，grandparent 更新为红色
cur->_col = BLACK;
grandparent->_col = RED;
}
//旋转之后，不需要继续更新了
break;
}
}
}
//根节点可能被更新为红色，要重新变为黑色
_root->_col = BLACK;
return { Iterator(newnode, _root), true};
}
//将 Find 返回值修改为 Iterator
Iterator 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 Iterator(cur, _root);
}
return End();
}
//...
private:
Node* _root = nullptr;
};
//mysey.hpp
#pragma once
#include "RBTree.hpp"
namespace LTL
{
template<class K>
class set
{
public:
//...
pair<iterator, bool> insert(const K& key)
{
return _t.Insert(key);
}
iterator find(const K& key)
{
return _t.Find(key);
}
private:
//将 K 修改为 const 即可
//这里修改了，别忘记上面 typedef 部分也要修改
RBTree<K, const K, SetKeyOfT> _t;
};
}
//mymap.hpp
#pragma once
#include "RBTree.hpp"
namespace LTL
{
template<class K, class V>
class map
{
public:
//...
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 = insert({ key, V() });
return ret.first->second;
}
private:
//只将 key 修改为 const 即可
RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};
}完整以及测试代码
cpp
//RBTree.hpp
#pragma once
#include <iostream>
using namespace std;
enum Color
{
BLACK,
RED
};
//节点的类型也要修改为 T
template<class T>
struct RBNode
{
T _data;
RBNode<T>* _left;
RBNode<T>* _right;
RBNode<T>* _parent;
Color _col;
RBNode(const T& data)
:_data(data)
,_left(nullptr)
, _right(nullptr)
,_parent(nullptr)
,_col(BLACK)
{}
};
//与 list 相同，Ptr、Ref 模板参数主要是方便实现 iterator 或者 const_iterator 的
template<class T, class Ptr, class Ref>
struct _Rb_Iterator
{
typedef RBNode<T> Node;
typedef _Rb_Iterator<T, Ptr, Ref> Self;
Node* _node;
//这里有一个 _root 主要是方便 end() 做--
Node* _root;
_Rb_Iterator(Node* node, Node* root)
:_node(node)
,_root(root)
{}
Ref operator*()
{
return _node->_data;
}
Ptr operator->()
{
return &(_node->_data);
}
bool operator!=(const Self& it)
{
return _node != it._node;
}
bool operator==(const Self& it)
{
return _node == it._node;
}
//前置++
Self& operator++()
{
//1. 如果当前节点的右孩子不为空
if (_node->_right)
{
//去到右孩子的最左节点
_node = _node->_right;
while (_node->_left)
{
_node = _node->_left;
}
}
//2. 右孩子为空
else
{
//去到当前节点为父亲节点左孩子的父亲节点
Node* cur = _node;
Node* parent = cur->_parent;
//parent 可能为空，比如图中的 50 节点
while (parent && cur == parent->_right)
{
cur = parent;
parent = parent->_parent;
}
_node = parent;
}
return *this;
}
//后置++
Self operator++(int)
{
Self tmp = *this;
++(*this);
return tmp;
}
//前置--
Self& operator--()
{
//operator--中可以对 end() 进行--
if (_node == nullptr)
{
_node = _root;
return *this;
}
//与 operator++ 的逻辑相反
//1. 左孩子不为空
if (_node->_left)
{
//去到左孩子的最右节点
_node = _node->_left;
while (_node->_right)
{
_node = _node->_right;
}
}
//2. 左孩子为空
else
{
//去到当前节点为父亲节点右孩子的父亲节点
Node* cur = _node;
Node* parent = cur->_parent;
//parent 可能为空，比如图中的 50 节点
while (parent && cur == parent->_left)
{
cur = parent;
parent = parent->_parent;
}
_node = parent;
}
return *this;
}
//后置--
Self operator--(int)
{
Self tmp = *this;
--(*this);
return tmp;
}
};
//K -- key，T -- key/pair<K, V>
template<class K, class T, class KeyOfT>
class RBTree
{
typedef RBNode<T> Node;
void RotateR(Node* parent)
{
Node* Parent = parent->_parent;
Node* subLR = parent->_left->_right;
Node* subL = parent->_left;
//不要忘记更改父亲的指向
//parent 可能为根节点
if (parent == _root)
_root = subL;
else
{
if (parent == Parent->_left)
Parent->_left = subL;
else
Parent->_right = subL;
}
//subL 的父亲节点也需要更改
subL->_parent = Parent;
//更改左右孩子指向
parent->_left = subLR;
subL->_right = parent;
//不要忘记更改 parent 以及 subLR 的父亲指向
parent->_parent = subL;
if (subLR)
subLR->_parent = parent;
}
void RotateL(Node* parent)
{
Node* Parent = parent->_parent;
Node* subR = parent->_right;
Node* subRL = subR->_left;
//不要忘记更改父亲的指向
//parent 可能为根节点
if (parent == _root)
_root = subR;
else
{
if (parent == Parent->_left)
Parent->_left = subR;
else
Parent->_right = subR;
}
//subR 的父亲节点也需要更改
subR->_parent = Parent;
//更改孩子指向
subR->_left = parent;
parent->_right = subRL;
//不要忘记更改 parent 与 subRL 的父亲指向
parent->_parent = subR;
if (subRL)
subRL->_parent = parent;
}
void RotateLR(Node* parent)
{
Node* subL = parent->_left;
Node* subLR = subL->_right;
//先以 subL 为 parent 进行左单旋
RotateL(subL);
//再以 parent 为 parent 进行右单旋
RotateR(parent);
}
void RotateRL(Node* parent)
{
Node* subR = parent->_right;
Node* subRL = subR->_left;
//先以 subR 为 parent 进行右单旋
RotateR(subR);
//再以 parent 为 parent 进行左单旋
RotateL(parent);
}
bool Check(Node* root, int blackNum, const int refNum)
{
//前序遍历走到空，说明一条路径走完了
if (root == nullptr)
{
if (blackNum != refNum)
{
cout << "路径黑色节点不同" << endl;
return false;
}
return true;
}
//检查其与其父亲节点是否都是红色
if (root->_col == RED && root->_parent && root->_parent->_col == RED)
{
cout << "存在连续的红色节点" << endl;
return false;
}
if (root->_col == BLACK)
++blackNum;
//如果左右子树都是红黑树，那当前子树就是红黑树
return Check(root->_left, blackNum, refNum) && Check(root->_right, blackNum, refNum);
}
int _Height(Node* root) const
{
if (root == nullptr)
return 0;
int LeftHeight = _Height(root->_left);
int RightHeight = _Height(root->_right);
return LeftHeight > RightHeight ? LeftHeight + 1 : RightHeight + 1;
}
size_t _size(Node* root) const
{
//利用递归来算节点个数
if (root == nullptr)
return 0;
return _size(root->_left) + _size(root->_right) + 1;
}
public:
typedef _Rb_Iterator<T, T*, T&> Iterator;
typedef _Rb_Iterator<T, const T*, const T&> Const_Iterator;
Iterator Begin()
{
//找到最左节点
Node* cur = _root;
while (cur->_left)
{
cur = cur->_left;
}
return Iterator(cur, _root);
}
Iterator End()
{
return Iterator(nullptr, _root);
}
Const_Iterator Begin() const
{
//找到最左节点
Node* cur = _root;
while (cur->_left)
{
cur = cur->_left;
}
return Const_Iterator(cur, _root);
}
Const_Iterator End() const
{
return Const_Iterator(nullptr, _root);
}
pair<Iterator, bool> Insert(const T& data)
{
if (_root == nullptr)
{
_root = new Node(data);
return { Iterator(_root, _root), true };
}
Node* cur = _root;
Node* parent = nullptr;
KeyOfT kot;
while (cur)
{
if (kot(cur->_data) > kot(data))
{
parent = cur;
cur = cur->_left;
}
else if (kot(cur->_data) < kot(data))
{
parent = cur;
cur = cur->_right;
}
else
return { Iterator(cur, _root), false };
}
cur = new Node(data);
//保存一下新插入的节点
Node* newnode = cur;
if (kot(parent->_data) > kot(data))
parent->_left = cur;
else
parent->_right = cur;
//不要忘记链接 _parent
cur->_parent = parent;
//新插入节点的颜色为红色
cur->_col = RED;
//父亲不为空，并且 parent 颜色为红色时才需要更新颜色
while (parent && parent->_col == RED)
{
//parent 为红色，grandparent 一定存在
Node* grandparent = parent->_parent;
if (parent == grandparent->_left)
{
Node* uncle = grandparent->_right;
//如果 uncle 存在且为红色
if (uncle && uncle->_col == RED)
{
//将 uncle 与 parent 变黑，grandparent 变红
uncle->_col = parent->_col = BLACK;
grandparent->_col = RED;
//继续向上更新
cur = grandparent;
parent = cur->_parent;
}
else
{
//uncle 不存在或者 uncle 存在且为黑
if (cur == parent->_left)
{
//不管哪种情况，只要在左边，就进行右单旋
RotateR(grandparent);
//将 parent 更新为黑色，grandparent 更新为红色
parent->_col = BLACK;
grandparent->_col = RED;
}
else
{
//进行左右双旋
RotateLR(grandparent);
//将 cur 更新为黑色，grandparent 更新为红色
cur->_col = BLACK;
grandparent->_col = RED;
}
//旋转之后，不需要更新了
break;
}
}
else
{
Node* uncle = grandparent->_left;
//如果 uncle 存在且为红色
if (uncle && uncle->_col == RED)
{
//将 uncle 与 parent 变黑，grandparent 变红
uncle->_col = parent->_col = BLACK;
grandparent->_col = RED;
//继续向上更新
cur = grandparent;
parent = cur->_parent;
}
else
{
//uncle 不存在或者 uncle 存在且为黑
if (cur == parent->_right)
{
//不管哪种情况，只要在右边，就进行左单旋
RotateL(grandparent);
//将 parent 更新为黑色，grandparent 更新为红色
parent->_col = BLACK;
grandparent->_col = RED;
}
else
{
//进行右左双旋
RotateRL(grandparent);
//将 cur 更新为黑色，grandparent 更新为红色
cur->_col = BLACK;
grandparent->_col = RED;
}
//旋转之后，不需要继续更新了
break;
}
}
}
//根节点可能被更新为红色，要重新变为黑色
_root->_col = BLACK;
return { Iterator(newnode, _root), true};
}
Iterator 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 Iterator(cur, _root);
}
return End();
}
bool IsRBTree()
{
//如果是一棵空树，那就是一棵红黑树
if (_root == nullptr)
return true;
if (_root->_col == RED)
{
cout << "根节点颜色为红色" << endl;
return false;
}
//计算出最左边路径的黑色节点个数
int blackNum = 0;
Node* cur = _root;
while (cur)
{
if (cur->_col == BLACK)
blackNum++;
cur = cur->_left;
}
return Check(_root, 0, blackNum);
}
size_t Height() const
{
return _Height(_root);
}
size_t size() const
{
return _size(_root);
}
private:
Node* _root = nullptr;
};
//myset.hpp
#pragma once
#include "RBTree.hpp"
namespace LTL
{
template<class K>
class set
{
public:
struct SetKeyOfT
{
const K& operator()(const K& key)
{
return key;
}
};
//实现迭代器
//这里需要在前面加 typename 关键字表明 Iterator 是个类型
typedef typename RBTree<K, const K, SetKeyOfT>::Iterator iterator;
typedef typename RBTree<K, const K, SetKeyOfT>::Const_Iterator 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:
//将 K 修改为 const 即可
//这里修改了，别忘记上面 typedef 部分也要修改
RBTree<K, const K, SetKeyOfT> _t;
};
}
//mymap.hpp
#pragma once
#include "RBTree.hpp"
namespace LTL
{
template<class K, class V>
class map
{
public:
struct MapKeyOfT
{
const K& operator()(const pair<K, V>& kv)
{
return kv.first;
}
};
//这里需要在前面加 typename 关键字表明 Iterator 是个类型
typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::Iterator iterator;
typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::Const_Iterator 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 = insert({ key, V() });
return ret.first->second;
}
private:
//只将 key 修改为 const 即可
RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};
}
//测试 Main.cc 代码
#include <iostream>
#include "myset.hpp"
#include "mymap.hpp"
#include <string>
using namespace std;
void TestSet()
{
LTL::set<int> s;
//测试插入
s.insert(1);
s.insert(10);
s.insert(6);
s.insert(-1);
s.insert(2);
//测试迭代器
auto it = s.begin();
while (it != s.end())
{
cout << *it << ' ';
++it;
}
cout << endl;
for (auto& e : s)
{
// key 不可修改
//++e;
cout << e << ' ';
}
cout << endl;
//测试查找
auto pos = s.find(10);
cout << *pos << endl;
pos = s.find(20);
if (pos == s.end())
{
cout << "end()" << endl;
}
}
void TestMap()
{
LTL::map<string, string> m;;
//测试插入
m.insert({ "sort", "排序" });
m.insert({ "left", "左边" });
m.insert({ "right", "右边" });
m.insert({ "algorithm", "算法" });
m.insert({ "get", "得到" });
//测试迭代器
for (auto& e : m)
{
cout << e.first << ":" << e.second << endl;
}
cout << endl;
//测试 operator[]
m["insert"] = "插入";
m["get"] = "获取";
for (auto& e : m)
{
cout << e.first << ":" << e.second << endl;
}
cout << endl;
}
int main()
{
//TestSet();
TestMap();
return 0;
}3 总结
其实 map 与 set 封装的核心并不是实现 map 与 set 本身，而是实现其底层的红黑树本身，map 和 set 其实就是在 RBTree 的基础上封装了一下而已。
为了一棵树实现两种容器，所以我们为 RBTree 中添加了 K 这个模板参数，T 才是树中真正存储的数据类型。但是这样在 insert 与 find 时，我们无法获取 data 的 key 值，所以我们引入了 KeyOfT 这个模板参数来获取 data 的 key 值，上层的 set 和 map 再传入获取 key 的仿函数，这样就可以进行 data 的比较了。
map 和 set 的迭代器，其实就是红黑树中序遍历的过程。前向中序遍历一共分为两种情况，一种是右子树不为空的情况，此时直接走到右子树的最左节点；另一种就是右子树为空的情况，此时走到孩子节点为其父亲节点左孩子的父亲节点处。而后向中序遍历正好与前向中序遍历相反，一种是左子树不为空的情况，此时直接走到左子树的最右节点；另一种就是左子树为空，此时走到孩子节点为其父亲节点右孩子的父亲节点处。
最后，我们要添加 key 不能修改的情况，也就是为模板参数中的 K 添加 const 即可。
总之，只要实现了对应的 RBTree，那么 set 和 map 的封装其实就很简单了，只是需要增加哎一些设计模块而已。