C++ 红黑树模拟实现 STL 中的 set 和 map
1. 分析源码
map 和 set 在 STL 库中是由红黑树来实现的。我们使用上一章所讲过的红黑树来模拟实现'阉割版'的 map 和 set,目的是加深对红黑树和这些容器的了解。
基于红黑树数据结构模拟实现 C++ STL 容器 set 和 map。通过分析 SGI-STL 源码结构,调整 rb_tree 泛型参数以支持键值对及单一键搜索场景。实现自定义 KeyOfT 仿函数提取比较键,封装双向迭代器支持中序遍历。重点阐述插入平衡逻辑、旋转操作及迭代器 ++/-- 的实现细节,最终复用红黑树构建 set 去重容器与 map 映射容器,强化对底层数据结构的理解。

map 和 set 在 STL 库中是由红黑树来实现的。我们使用上一章所讲过的红黑树来模拟实现'阉割版'的 map 和 set,目的是加深对红黑树和这些容器的了解。
SGI-STL30 版本源代码中,map 和 set 的源代码在 stl_map.h/stl_set.h/stl_tree.h 等头文件中。
核心结构如下:
// stl_set.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; // red-black tree representing set
};
// stl_map.h
template<class Key, class T, class Compare = less<Key>, class Alloc = alloc>
class map {
public:
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
};
通过对源码的分析,发现 rb_tree 使用了泛型思想。第二个模板参数 Value 决定了 _rb_tree_node 中存储的数据类型。
rb_tree 时第二个模板参数给的是 key。rb_tree 时第二个模板参数给的是 pair<const key, T>。这样一棵红黑树既可以实现 key 搜索场景的 set,也可以实现 key/value 搜索场景的 map。
注意源码里第一个模板参数 Key 是传给 find/erase 等函数做形参的类型的。对于 set 而言两个参数是一样的,但是对于 map 而言就不一样了,map insert 的是 pair 对象,但是 find 和 erase 的是 Key 对象。
我们需要对红黑树进行更改,使其支持泛型。相比源码调整一下,key 参数用 K,value 参数用 V,红黑树中的数据类型使用 T。
因为 RBTree 实现了泛型不知道 T 参数到底是 K,还是 pair<K, V>,所以 insert 内部进行比较时,需要通过仿函数把 key 取出来。
对 map 而言,MapKeyOfT 直接返回 pair<key, value> 中的 key:
struct MapKeyOfT {
const K& operator()(const pair<K, V>& kv) {
return kv.first;
}
};
对 set 而言,SetKeyOfT 直接返回 key:
struct SetKeyOfT {
const K& operator()(const K& key) {
return key;
}
};
红黑树的参数类型变为:
template<class K, class T, class KeyOfT>
class RBTree {
// ...
};
set 和 map 的迭代器都是双向迭代器,通过封装底层红黑树的迭代器来实现。
红黑树的 iterator 实现的框架跟 list 的 iterator 思路一致,用一个类型封装结点的指针,再通过重载运算符实现。
最重要的是 operator++ 和 operator-- 的实现。map 和 set 的迭代器走的是中序遍历(左子树->根结点->右子树)。
it 指向的结点的右子树不为空,找右子树的最左结点。it 指向的结点的右子树为空,沿着当前结点到根的祖先路径向上找,直到找到孩子是父亲左的那个祖先。如果没有找到孩子是父亲左的那个祖先,这时父亲为空了,也就是根节点的父亲,我们就把 it 中的结点指针置为 nullptr,用 nullptr 去充当 end()。
需要注意的是在 STL 源码中,红黑树增加了一个哨兵位头结点作为 end(),而我们用 nullptr 作为 end(),差别不大。
iterator 不支持修改,我们把 set 的第二个模板参数改成 const K 即可 RBTree<K, const K, SetKeyOfT> _titerator 不支持修改 key 但是可以修改 value,我们把 map 的第二个模板参数 pair 的第一个参数改成 const K 即可 RBTree<K, pair<const K, V>, MapKeyOfT> _tmap 要支持 [] 主要需要修改 insert 返回值支持,修改 RBtree 中的 insert 返回值为 pair<Iterator, bool> Insert(const T& data)。
下面是修改后的红黑树的代码实现:
#pragma once
#include <iostream>
using namespace std;
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 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; }
Self& operator++() {
if (_node->_right) {
Node* minleft = _node->_right;
while (minleft->_left) {
minleft = minleft->_left;
}
_node = minleft;
} else {
Node* cur = _node;
Node* parent = cur->_parent;
while (parent && cur == parent->_right) {
cur = parent;
parent = parent->_parent;
}
_node = parent;
}
return *this;
}
Self& operator--() {
if (_node == nullptr) {
Node* minright = _root;
while (minright && minright->_right) {
minright = minright->_right;
}
_node = minright;
} else if (_node->_left) {
Node* minright = _node->_left;
while (minright->_right) {
minright = minright->_right;
}
_node = minright;
} else {
Node* cur = _node;
Node* parent = cur->_parent;
while (parent && cur == parent->_left) {
cur = parent;
parent = parent->_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>
class RBTree {
typedef RBTreeNode<T> Node;
public:
typedef RBTreeIterator<T, T&, T*> Iterator;
typedef RBTreeIterator<T, const T&, const T*> Const_Iterator;
Iterator begin() {
Node* cur = _root;
while (cur && cur->_left) {
cur = cur->_left;
}
return Iterator(cur, _root);
}
Iterator end() {
return Iterator(nullptr, _root);
}
Const_Iterator begin() const {
Node* cur = _root;
while (cur && cur->_left) {
cur = cur->_left;
}
return Iterator(cur, _root);
}
Const_Iterator end() const {
return Iterator(nullptr, _root);
}
pair<Iterator, bool> Insert(const T& data) {
if (_root == nullptr) {
_root = new Node(data);
_root->_col = BLACK;
return { Iterator(_root, _root), true };
}
KeyOfT kot;
Node* cur = _root;
Node* parent = nullptr;
while (cur) {
if (kot(cur->_data) < kot(data)) {
parent = cur;
cur = parent->_right;
} else if (kot(cur->_data) > kot(data)) {
parent = cur;
cur = parent->_left;
} else {
return { Iterator(cur, _root), false };
}
}
cur = new Node(data);
Node* newnode = cur;
cur->_col = RED;
if (kot(parent->_data) > kot(data)) {
parent->_left = cur;
} else {
parent->_right = cur;
}
cur->_parent = parent;
while (parent && parent->_col == RED) {
Node* grandfather = parent->_parent;
if (parent == grandfather->_left) {
Node* uncle = grandfather->_right;
if (uncle && uncle->_col == RED) {
parent->_col = BLACK;
uncle->_col = BLACK;
grandfather->_col = RED;
cur = grandfather;
parent = cur->_parent;
} else {
if (cur == parent->_left) {
RotateR(grandfather);
parent->_col = BLACK;
grandfather->_col = RED;
} else {
RotateL(parent);
RotateR(grandfather);
cur->_col = BLACK;
grandfather->_col = RED;
}
break;
}
} else {
Node* uncle = grandfather->_left;
if (uncle && uncle->_col == RED) {
parent->_col = BLACK;
uncle->_col = BLACK;
grandfather->_col = RED;
cur = grandfather;
parent = cur->_parent;
} else {
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 { 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 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) {
if (refNum != blackNum) {
cout << "存在黑色结点的数量不相等的路径" << endl;
return false;
}
return true;
}
if (root->_col == RED && root->_parent->_col == RED) {
cout << "存在连续的红色结点" << endl;
return false;
}
if (root->_col == BLACK) {
blackNum++;
}
return Check(root->_left, blackNum, refNum) && Check(root->_right, blackNum, refNum);
}
void RotateR(Node* parent) {
Node* subL = parent->_left;
Node* subLR = subL->_right;
Node* Pparent = parent->_parent;
parent->_left = subLR;
if (subLR) subLR->_parent = parent;
subL->_right = parent;
parent->_parent = subL;
if (parent == _root) {
_root = subL;
subL->_parent = nullptr;
} else {
if (parent == Pparent->_left) {
Pparent->_left = subL;
} else {
Pparent->_right = subL;
}
subL->_parent = Pparent;
}
}
void RotateL(Node* parent) {
Node* subR = parent->_right;
Node* subRL = subR->_left;
Node* Pparent = parent->_parent;
parent->_right = subRL;
if (subRL) subRL->_parent = parent;
subR->_left = parent;
parent->_parent = subR;
if (parent == _root) {
_root = subR;
subR->_parent = nullptr;
} else {
if (parent == Pparent->_left) {
Pparent->_left = subR;
} else {
Pparent->_right = subR;
}
subR->_parent = Pparent;
}
}
Node* _root = nullptr;
};
其实当我们修改完红黑树后,实现 map 和 set 就非常简单了,我们只需要复用红黑树即可。
#pragma once
#include "RBTree.h"
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>::Const_Iterator const_iterator;
iterator begin() { return _rbtree.begin(); }
iterator end() { return _rbtree.end(); }
const_iterator begin() const { return _rbtree.begin(); }
const_iterator end() const { return _rbtree.end(); }
pair<iterator, bool> insert(const K& key) { return _rbtree.Insert(key); }
iterator find(const K& key) { return _rbtree.Find(key); }
private:
RBTree<K, const K, SetKeyOfT> _rbtree;
};
#pragma once
#include "RBTree.h"
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>::Const_Iterator const_iterator;
iterator begin() { return _rbtree.begin(); }
iterator end() { return _rbtree.end(); }
const_iterator begin() const { return _rbtree.begin(); }
const_iterator end() const { return _rbtree.end(); }
pair<iterator, bool> insert(const pair<K, V>& kv) { return _rbtree.Insert(kv); }
V& operator[](const K& key) {
pair<iterator, bool> ret = _rbtree.Insert({ key, V() });
return ret.first->second;
}
iterator find(const K& key) { return _rbtree.Find(key); }
private:
RBTree<K, pair<const K, V>, MapKeyOfT> _rbtree;
};
我们通过自己实现红黑树并且复用它来构建简单的 map 和 set 可以帮助我们:
红黑树是一种复杂的自平衡二叉搜索树,其核心在于通过颜色规则(红 / 黑)和旋转操作维持平衡性,保证插入、删除、查找的时间复杂度稳定在 O(log n)。我们在实现红黑树的过程中,能深入理解:
而复用红黑树实现 map 和 set 时,能进一步理解这两种容器的底层原理:
红黑树是 map 和 set 的 '底层引擎'—— 两者都依赖红黑树的排序和查找能力,仅在节点存储的数据(set 存键,map 存键值对)和接口设计上有差异。通过复用红黑树实现这两种容器,能直观体会 '抽象复用' 的思想:

微信公众号「极客日志」,在微信中扫描左侧二维码关注。展示文案:极客日志 zeeklog
使用加密算法(如AES、TripleDES、Rabbit或RC4)加密和解密文本明文。 在线工具,加密/解密文本在线工具,online
将字符串编码和解码为其 Base64 格式表示形式即可。 在线工具,Base64 字符串编码/解码在线工具,online
将字符串、文件或图像转换为其 Base64 表示形式。 在线工具,Base64 文件转换器在线工具,online
将 Markdown(GFM)转为 HTML 片段,浏览器内 marked 解析;与 HTML转Markdown 互为补充。 在线工具,Markdown转HTML在线工具,online
将 HTML 片段转为 GitHub Flavored Markdown,支持标题、列表、链接、代码块与表格等;浏览器内处理,可链接预填。 在线工具,HTML转Markdown在线工具,online
通过删除不必要的空白来缩小和压缩JSON。 在线工具,JSON 压缩在线工具,online