跳到主要内容红黑树封装 map 和 set 的实现原理与代码 | 极客日志C++算法
红黑树封装 map 和 set 的实现原理与代码
本文基于 SGI STL 源码分析，讲解如何使用红黑树封装实现 C++ 标准库中的 map 和 set 容器。内容涵盖红黑树泛型设计、KeyOfT 仿函数作用、迭代器中序遍历逻辑（operator++/--）、map 的 operator[] 重载以及构造析构函数的实现。通过代码示例展示了核心功能的模拟实现过程。
一、了解源码
在 SGI STL30 版本中，看一下 map 和 set 的部分源码
#ifndef __SGI_STL_INTERNAL_TREE_H
#include <stl_tree.h>
#endif
#include <stl_set.h>
#include <stl_multiset.h>
#ifndef __SGI_STL_INTERNAL_TREE_H
#include <stl_tree.h>
#endif
#include <stl_map.h>
#include <stl_multimap.h>
template<class Key, class Compare = less<Key>, class Alloc = alloc>
class set {
public:
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;
};
template<class Key, class , = less<Key>, Alloc = alloc>
map {
:
Key key_type;
T mapped_type;
pair< Key, T> value_type;
:
rb_tree<key_type, value_type, select1st<value_type>, key_compare, Alloc> rep_type;
rep_type t;
};
{
__rb_tree_color_type color_type;
__rb_tree_node_base* base_ptr;
color_type color;
base_ptr parent;
base_ptr left;
base_ptr right;
};
< , , , , = alloc>
rb_tree {
:
* void_pointer;
__rb_tree_node_base* base_ptr;
__rb_tree_node<Value> rb_tree_node;
rb_tree_node* link_type;
Key key_type;
Value value_type;
:
:
size_type node_count;
link_type header;
};
< >
: __rb_tree_node_base {
__rb_tree_node<Value>* link_type;
Value value_field;
};
微信扫一扫，关注极客日志
微信公众号「极客日志」，在微信中扫描左侧二维码关注。展示文案：极客日志 zeeklog
相关免费在线工具
- 加密/解密文本
使用加密算法（如AES、TripleDES、Rabbit或RC4）加密和解密文本明文。 在线工具，加密/解密文本在线工具，online
- Base64 字符串编码/解码
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具，Base64 字符串编码/解码在线工具，online
- Base64 文件转换器
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具，Base64 文件转换器在线工具，online
- Markdown 转 HTML
将 Markdown（GFM）转为 HTML 片段，浏览器内 marked 解析；与 HTML 转 Markdown 互为补充。 在线工具，Markdown 转 HTML在线工具，online
- HTML 转 Markdown
将 HTML 片段转为 GitHub Flavored Markdown，支持标题、列表、链接、代码块与表格等；浏览器内处理，可链接预填。 在线工具，HTML 转 Markdown在线工具，online
- JSON 压缩
通过删除不必要的空白来缩小和压缩JSON。 在线工具，JSON 压缩在线工具，online
T
class
Compare
class
class
public
typedef
typedef
typedef
const
private
typedef
struct
__rb_tree_node_base
typedef
typedef
template
class
Key
class
Value
class
KeyOfValue
class
Compare
class
Alloc
class
protected
typedef
void
typedef
typedef
typedef
typedef
typedef
public
protected
template
class
Value
struct
__rb_tree_node
public
typedef
部分源码如上，我们通过源码可以看到源码中 rb_tree 使用了泛型思维实现；其中 rb_tree 是实现 key 搜索场景还是实现 key/value 的搜索场景不是写死的，而是通过第二个模板参数来决定的。
set 在实例化时第二个模板参数给的是 key，而 map 实例化时第二个模板参数给的是 pair<const K, T>，这样我们一颗红黑树就可以实现两种效果，既可以实现 key 搜索场景的 set 也可以实现 key/value 搜索场景 map。这里源码当中，模板参数用 T 来代替了 value，而其中 value_type 并不是我们之前 key/value 场景中的 value；源码当中 value_type 反而是红黑树节点中存储的数据类型。（这样 set 存储的数据就是 K，而 map 中存储的数据就是 pair<K , V>）。这里第二个模板参数 Value 已经控制了红黑树节点中存储的数据类型，那为什么还需要第一个模板参数呢？，就比如 set 模板参数传的是 K,K；这里主要原因还是因为 map 和 set 中 find 和 erase 函数的参数都是 Key,所以需要第一个模板参数传给 find 和 erase。对于 set 而言两个参数都是 key，而对于 map 而言，map 的 insert 是 pair 对象，但是 find 和 erase 的是 Key 对象。
二、封装实现 map 和 set
我们复用之前实现的 红黑树 来实现 set 和 map；我们需要进行一系列的修改，让我们实现的 红黑树 可以完成复用，让我们实现出来 set 和 map。
要复用整个框架，我们首先要进行一些修改；这里对比源码当中实现方法进行一系列修改和完善：
模板参数：红黑树的模板参数要用三个 <class K , class T , class KeyOfT>，能够进行 insert 操作。迭代器：RBTree 要实现出迭代器，能够进行遍历。map 中的 []：对于 map，我们要实现 [] 运算符重载；那个实现其完整的 [] 操作。
1. 复用红黑树框架，实现 insert
首先来看如何去封装，实现 map 和 set 的框架。
对于三个模板参数 K , T , KeyOfT，各有什么用呢？
key 和 T 就不多解释了，在了解源码时已经知道了它们的作用。
KeyOfT 的作用是什么？
首先我们知道，红黑树我们修改成 K , T 结构之后，我们在 insert 的时候函数参数是 T，我们并不知道 T 是 K 还是 pair<K , V>；那我们如何去进行大小比较呢？
这里我们要看一下 pair<K , V> 的比较方法是什么
可以看到库里面 pair 的比较方法是先比较 first，first 大就大，first 小就小；如果相等那就比较 second。
我们不想要这样的比较方法，我们就想要单纯的比较 K 也就是 first。（这里我们想要的比较方法，关键我们需要拿到 K）
KeyOfT 是一个仿函数，它重载了 operator() 方法，它的作用就是返回 T 中的 K；
如果 T 是 K 那就直接返回，如果 T 是 pair<K , V> 就返回其中的 K。
那我们在 RBTree 怎么知道 T 要返回的是什么呢？（显然是无法知道的，但是上层复用的肯定知道；所以我们就在 set 和 map 中写好 SetKeyOfT 和 MapKeyOfT，将其放在第三个模板参数的位置；这样我们在下层 RBTree 中就能够拿到 T 中的 K 了。
支持 insert 实现
现在就对于上述来修改我们之前实现好的 RBTree 代码。
enum Color { RED, BLACK };
template<class T>
struct RBTreeNode {
RBTreeNode* _left;
RBTreeNode* _right;
RBTreeNode* _parent;
T _data;
Color _col;
RBTreeNode(const T& data) :_data(data),_left(nullptr),_right(nullptr),_parent(nullptr),_col(RED) {}
};
template<class K, class T, class KOT>
class RBTree {
typedef RBTreeNode<T> Node;
public:
bool insert(const T& data) {
if (_root == nullptr)
{
_root = new Node(data);
_root->_col = BLACK;
return true;
}
Node* tail = _root;
Node* parent = nullptr;
KOT kot;
while (tail) {
if (kot(data) > kot(tail->_data)) {
parent = tail;
tail = tail->_right;
}
else if (kot(data) < kot(tail->_data)) {
parent = tail;
tail = tail->_left;
}
else {
return false;
}
}
Node* cur = new Node(data);
cur->_col = RED;
cur->_parent = parent;
if (kot(cur->_data) > kot(parent->_data)) {
parent->_right = cur;
}
else if (kot(cur->_data) < kot(parent->_data)) {
parent->_left = cur;
}
while (parent && parent->_col == RED) {
Node* grandfather = parent->_parent;
if (parent == grandfather->_left) {
if (uncle && uncle->_col == RED)
{
uncle->_col = parent->_col = BLACK;
grandfather->_col = RED;
cur = grandfather;
parent = cur->_parent;
}
else {
if (cur == parent->_left)
{
RotateR(grandfather);
parent->_col = BLACK;
grandfather->_col = RED;
break;
}
else
{
RotateL(parent);
RotateR(grandfather);
cur->_col = BLACK;
grandfather->_col = RED;
break;
}
}
}
else {
Node* uncle = grandfather->_left;
if (uncle && uncle->_col == RED)
{
uncle->_col = parent->_col = BLACK;
grandfather->_col = RED;
cur = grandfather;
parent = cur->_parent;
}
else {
if (cur == parent->_right)
{
RotateL(grandfather);
parent->_col = BLACK;
grandfather->_col = RED;
break;
}
else
{
RotateR(parent);
RotateL(grandfather);
cur->_col = BLACK;
grandfather->_col = RED;
break;
}
}
}
}
_root->_col = BLACK;
return true;
}
private:
void RotateR(Node* parent)
{
Node* subl = parent->_left;
Node* sublr = parent->_left->_right;
parent->_left = sublr;
if (sublr) sublr->_parent = parent;
Node* ppNode = parent->_parent;
parent->_parent = subl;
subl->_parent = ppNode;
if (ppNode == nullptr) {
_root = subl;
}
else {
if (parent == ppNode->_left) {
ppNode->_left = subl;
}
else if (parent->_right) {
ppNode->_right = subl;
}
}
}
void RotateL(Node* parent)
{
Node* subr = parent->_right;
Node* subrl = parent->_right->_left;
parent->_right = subrl;
if (subrl) subrl->_parent = parent;
Node* ppNode = parent->_parent;
parent->_parent = subr;
subr->_left = parent;
subr->_parent = ppNode;
if (ppNode == nullptr) {
_root = subr;
}
else {
if (parent == ppNode->_left) {
ppNode->_left = subr;
}
else if (parent == ppNode->_right) {
ppNode->_right = subr;
}
}
}
Node* _root = nullptr;
};
namespace HL {
template<class K>
class set {
struct SetKeyOfT {
const K& operator()(const K& key) {
return key;
}
};
public:
set() {}
bool insert(const K& k) {
return _rb.insert(k);
}
private: RBTree<K, const K, SetKeyOfT> _rb;
};
}
namespace HL {
template<class K, class V>
class map {
struct MapKeyOfT {
const K& operator()(const pair<const K, V>& kv) {
return kv.first;
}
};
public:
map() {}
bool insert(const pair<const K, V>& kv) {
return _rb.insert(kv);
}
private: RBTree<K, pair<const K, V>, MapKeyOfT> _rb;
};
}
2. 实现 iterator，完成 set 和 map 的迭代器遍历
对于 RBTree 的迭代器，它和 list 的迭代器是非常相似的，都是对于 节点的指针 进行封装来实现的；通过实现运算符重载来让迭代器能够像 指针 那样一样的访问行为。
运算符重载
这里先简单来实现 operator==、operator!=、operator*、operator-> 等这些简单一些的运算符重载
template<class T, class Ref, class Ptr>
struct RBTreeIterator {
typedef RBTreeNode<T> Node;
typedef RBTreeIterator<T, Ref, Ptr> Self;
Node* _node;
Node* _root;
RBTreeIterator(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++() {
if (_node->_right) {
Node* cur = _node->_right;
while (cur && cur->_left) {
cur = cur->_left;
}
_node = cur;
}
else {
Node* parent = _node->_parent;
Node* cur = _node;
while (parent && cur == parent->_right) {
cur = parent;
parent = cur->_parent;
}
_node = parent;
}
return *this;
}
Self& operator--() {
if (_node == nullptr) {
Node* cur = _root;
while (cur && cur->_right) {
cur = cur->_right;
}
_node = cur;
}
else {
if (_node->_left) {
Node* cur = _node->_left;
while (cur && cur->_right) {
cur = cur->_right;
}
_node = cur;
}
else {
Node* parent = _node->_parent;
Node* cur = _node;
while (parent && cur == parent->_left) {
cur = parent;
parent = cur->_parent;
}
_node = parent;
}
}
return *this;
}
};
RBTree、set 和 map 的迭代器实现
这里 begin() 返回的是中序遍历的第一个位置，也就是红黑树的最左节点；而 end() 返回的则是 nullptr 位置
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);
}
ConstIterator begin() const {
Node* cur = _root;
while (cur && cur->_left) {
cur = cur->_left;
}
return Iterator(cur, _root);
}
ConstIterator end() const {
return Iterator(nullptr, _root);
}
这里 set 和 map 迭代器实现起来就简单很多了，直接对 RBTree 的迭代器进行复用即可，而 begin() 和 end() 就直接进行调用即可。
typedef typename RBTree<K, const K, SetKeyOfT>::Iterator iterator;
typedef typename RBTree<K, const K, SetKeyOfT>::ConstIterator const_iterator;
iterator begin() {
return _rb.begin();
}
iterator end() {
return _rb.end();
}
const_iterator begin() const {
return _rb.begin();
}
const_iterator end() const {
return _rb.end();
}
这里 typedef typename RBTree<K, K, SetKeyOfT>::Iterator iterator; 中 typename 是告诉编译器这是个类型，让编译器在类模版实例化之后去取其中的类型名（因为编译器是按需实例化的，不加可能会报错）C++ 语法是运行这样的。
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 _rb.begin();
}
iterator end() {
return _rb.end();
}
const_iterator begin() const {
return _rb.begin();
}
const_iterator end() const {
return _rb.end();
}
测试 set 和 map
这里来测试一下实现的 set 和 map 的迭代器
void test_map() {
HL::map<int, int> mp;
mp.insert({ 1,1 });
mp.insert({ 2,2 });
mp.insert({ 3,3 });
mp.insert({ 4,4 });
mp.insert({ 5,5 });
auto it = mp.begin();
while (it != mp.end()) {
cout << it->first << ": " << it->second << endl;
++it;
}
cout << endl;
it = mp.end();
while (it != mp.begin()) {
--it;
cout << it->first << ": " << it->second << endl;
}
cout << endl;
}
void test_set() {
HL::set<int> st;
st.insert(1);
st.insert(2);
st.insert(3);
st.insert(4);
st.insert(5);
st.insert(6);
st.insert(7);
st.insert(8);
auto it = st.begin();
while (it != st.end()) {
cout << *it << " ";
++it;
}
cout << endl;
it = st.end();
while (it != st.begin()) {
--it;
cout << *it << " ";
}
cout << endl;
}
int main() {
test_set();
test_map();
return 0;
}
可以看到我们已经完成 set 和 map 的迭代器遍历。
三、map 的 operator[] 的实现
对于 map 的 operator[]，我们知道它有 两种 功能：
一是 key 存在，通过 key 查找 value，返回 value 的引用；
二是如果 key 不存在进插入 key，其对应的 value 是缺省值，然后返回 value 的引用；
所以这里不管 key 是否存在，我们都要返回对应的 value 的引用；
我们在实现时就肯定需要调用 insert 函数，先来看一下库里面 map 的 insert 函数
可以看到，相比于我们现在实现的 insert，库里面返回的是 pair<iterator , bool> 类型，这样我们就能拿到插入位置的迭代器。
所以我们在实现的时候也需要这样去设计；（再结合一下 insert 的逻辑，如果插入值存在，我们就可以返回当前位置迭代器和 false，这样就达到了我们想要的效果）
pair<Iterator, bool> insert(const T& data) {
if (_root == nullptr)
{
_root = new Node(data);
_root->_col = BLACK;
return make_pair(Iterator(_root, _root), true);
}
Node* tail = _root;
Node* parent = nullptr;
KOT kot;
while (tail) {
if (kot(data) > kot(tail->_data)) {
parent = tail;
tail = tail->_right;
}
else if (kot(data) < kot(tail->_data)) {
parent = tail;
tail = tail->_left;
}
else {
return make_pair(Iterator(tail, _root), false);
}
}
Node* cur = new Node(data);
Node* newnode = cur;
cur->_col = RED;
cur->_parent = parent;
if (kot(cur->_data) > kot(parent->_data)) {
parent->_right = cur;
}
else if (kot(cur->_data) < kot(parent->_data)) {
parent->_left = cur;
}
while (parent && parent->_col == RED) {
Node* grandfather = parent->_parent;
if (parent == grandfather->_left) {
if (uncle && uncle->_col == RED)
{
uncle->_col = parent->_col = BLACK;
grandfather->_col = RED;
cur = grandfather;
parent = cur->_parent;
}
else {
if (cur == parent->_left)
{
RotateR(grandfather);
parent->_col = BLACK;
grandfather->_col = RED;
break;
}
else
{
RotateL(parent);
RotateR(grandfather);
cur->_col = BLACK;
grandfather->_col = RED;
break;
}
}
}
else {
Node* uncle = grandfather->_left;
if (uncle && uncle->_col == RED)
{
uncle->_col = parent->_col = BLACK;
grandfather->_col = RED;
cur = grandfather;
parent = cur->_parent;
}
else {
if (cur == parent->_right)
{
RotateL(grandfather);
parent->_col = BLACK;
grandfather->_col = RED;
break;
}
else
{
RotateR(parent);
RotateL(grandfather);
cur->_col = BLACK;
grandfather->_col = RED;
break;
}
}
}
}
_root->_col = BLACK;
return make_pair(Iterator(newnode, _root), true);
}
pair<iterator, bool> insert(const pair<const K, V>& kv) {
return _rb.insert(kv);
}
V& operator[](const K& key) {
pair<iterator, bool> ret = insert({ key, V() });
return ret.first->second;
}
到这里我们实现的 map 就可以使用 [] 去访问了。
这里来测试一下我们实现的是否到达了我们预期的效果：
void test_map() {
HL::map<string, string> mp;
mp.insert({ "sort", "排序" });
mp.insert({ "left", "左边" });
mp.insert({ "right", "右边" });
mp["left"] = "左边，剩余";
mp["insert"] = "插入";
mp["string"];
HL::map<string, string>::iterator it = mp.begin();
while (it != mp.end()) {
it->second += 'x';
cout << it->first << ":" << it->second << endl;
++it;
}
cout << endl;
}
这里简单使用 map 中 []，是可以到达我们的预期效果的。
四、map 和 set 的构造函数与析构函数
对于构造函数，我们直接进行复用 insert 即可。
set() {}
template<class InputIterator>
set(InputIterator first, InputIterator last) {
while (first != last) {
_rb.insert(*first);
first++;
}
}
set(const set& x) {
for (const auto& e : x) {
_rb.insert(e);
}
}
map 和 set 一样直接调用 RBTree 中的 insert 即可。
map() {}
template<class InputIterator>
map(InputIterator first, InputIterator last) {
while (first != last) {
_rb.insert(*first);
first++;
}
}
map(const set& x) {
for (const auto& e : x) {
_rb.insert(e);
}
}
对于析构函数，我们就要实现 RBTree 的析构函数，然后在 map 和 set 中进行调用；这里就使用递归来释放每一个节点。
~RBTree() {
Destroy(_root);
}
void Destroy(Node* root) {
if (root == nullptr)
return;
Destroy(root->_left);
Destroy(root->_right);
delete root;
}
而 map 和 set 就直接 delete 即可，编译器会自动调用 RBTree 的析构函数。
五、总结
这里并没有去实现 erase 删除功能，只是简单模拟了 map 和 set 的实现。
首先就是一个红黑树框架来封装实现 map 和 set 其次就是模拟迭代器的实现，整个思路（重点：operator++ 和 operator--）然后就是 map 中 operator[] 的特殊功能。最后就是构造函数、析构函数等这些简单的复用。
这里并没有实现像 size、find、empty、count 等函数，这些就比较简单。
