跳到主要内容红黑树进阶:手撕 STL 源码封装 RB-tree 实现 map 和 set | 极客日志C++算法
红黑树进阶:手撕 STL 源码封装 RB-tree 实现 map 和 set
基于 C++ 泛型编程思想,通过模板参数适配红黑树结构,分别实现 key-only 的 set 与 key-value 的 map。核心涵盖节点定义、插入旋转逻辑、仿函数 KeyOfT 解决键值提取问题、迭代器中序遍历实现及 const 正确性保证。代码展示了从底层 rb_tree 到上层容器封装的完整流程,对比了与 SGI-STL 的设计差异,并验证了 operator[] 等接口功能。
黑客0 浏览 一。源码及框架分析
1.1 STL 源码中的设计思想
在 SGI-STL30 版本中,map 和 set 的实现基于一棵红黑树(rb_tree)。这种设计体现了 C++ 泛型编程的强大之处:通过同一棵红黑树,既可以实现 key-only 的 set,也可以实现 key/value 的 map。
核心设计思想:
- rb_tree 的节点中存储什么类型,由模板参数决定
- set 传给 rb_tree 的是 Key 类型
- map 传给 rb_tree 的是 pair<const Key, T> 类型
1.2 STL 源码框架分析
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 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;
};
{
__rb_tree_color_type color_type;
__rb_tree_node_base* base_ptr;
color_type color;
base_ptr parent;
base_ptr left;
base_ptr right;
};
< >
:public__rb_tree_node_base{
__rb_tree_node<Value>* link_type;
Value value_field;
};
struct
__rb_tree_node_base
typedef
typedef
template
class
Value
struct
__rb_tree_node
typedef
关键问题:为什么 rb_tree 需要两个模板参数 Key 和 Value?
- Value:决定了节点中存储的数据类型(set 是 Key,map 是 pair)
- Key:用于查找和删除操作的参数类型(find/erase 的参数类型)
对于 set 而言,Key 和 Value 是相同的;对于 map 而言,Key 是键的类型,Value 是 pair 类型。这种设计使得一棵红黑树既能处理 set 场景,也能处理 map 场景。
二。模拟实现 map 和 set(实现复用红黑树的框架)
2.1 红黑树节点的定义
#pragma once
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),_col(RED){}
};
2.2 红黑树的基本框架
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*> ConstIterator;
RBTree()=default;
~RBTree(){Destroy(_root); _root =nullptr;}
pair<Iterator,bool>Insert(const T& data);
Iterator Find(const K& key);
Iterator Begin(); Iterator End();
ConstIterator Begin()const; ConstIterator End()const;
private:
void RotateL(Node* parent);
void RotateR(Node* parent);
void Destroy(Node* root);
Node* _root =nullptr;
};
2.3 解决 Key 的比较问题:KeyOfT 仿函数
问题:红黑树在插入时需要比较键的大小,但 T 可能是 Key 类型,也可能是 pair 类型。pair 的比较规则是 first 和 second 一起比较,不符合我们的需求(我们只想比较 key)。
解决方案:使用仿函数 KeyOfT,从 T 对象中取出 Key 进行比较。
namespace bit {
template<class K,class V>
class map{
struct MapKeyOfT{
const K& operator()(const pair<K, V>& kv){return kv.first;
};
};
}
namespace bit {
template<class K>
class set{
struct SetKeyOfT{
const K& operator()(const K& key){return key;
};
};
}
2.4 支持 insert 插入
RBTree 中的 Insert 实现(关键部分是使用 KeyOfT 进行比较)
template<class K,class T,class KeyOfT>
pair<typename RBTree<K, T, KeyOfT>::Iterator,bool>RBTree<K, T, KeyOfT>::Insert(const T& data){
if(_root ==nullptr){
_root =new Node(data);
_root->_col = BLACK;
return make_pair(Iterator(_root, _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)){
parent = cur; cur = cur->_left;
}else{
return make_pair(Iterator(cur, _root),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;
while(parent && parent->_col == RED){
Node* grandfather = parent->_parent;
if(parent == grandfather->_left){
Node* uncle = grandfather->_right;
if(uncle && uncle->_col == RED){
parent->_col = 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 = 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 make_pair(Iterator(newnode, _root),true);
}
2.5 map 和 set 的 insert 封装
namespace bit {
template<class K,class V>
class map{
public:
pair<iterator,bool>insert(const pair<K, V>& kv){
return _t.Insert(kv);
}
private:
RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};
}
namespace bit {
template<class K>
class set{
public:
pair<iterator,bool>insert(const K& key){
return _t.Insert(key);
}
private:
RBTree<K,const K, SetKeyOfT> _t;
};
}
三。迭代器的实现
3.1 迭代器结构设计
operator++:前往中序下一个节点
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){}
Self& operator++();
Self& operator--();
Ref operator*(){return _node->_data;}
Ptr operator->(){return&_node->_data;}
bool operator!=(const Self& s)const{return _node != s._node;}
bool operator==(const Self& s)const{return _node == s._node;}
};
3.2 迭代器的++操作
template<class T,class Ref,class Ptr>
typename RBTreeIterator<T, Ref, Ptr>::Self& RBTreeIterator<T, Ref, Ptr>::operator++(){
if(_node ==nullptr) return this;
if(_node->_right){
Node* leftMost = _node->_right;
while(leftMost->_left){ leftMost = leftMost->_left;}
_node = leftMost;
}else{
Node* cur = _node;
Node* parent = cur->_parent;
while(parent && cur == parent->_right){ cur = parent; parent = cur->_parent;}
_node = parent;
}
return *this;
}
3.3 迭代器的–操作
template<class T,class Ref,class Ptr>
typename RBTreeIterator<T, Ref, Ptr>::Self& RBTreeIterator<T, Ref, Ptr>::operator--(){
if(_node ==nullptr){
Node* rightMost = _root;
while(rightMost && rightMost->_right){ rightMost = rightMost->_right;}
_node = rightMost;
}else if(_node->_left){
Node* rightMost = _node->_left;
while(rightMost->_right){ rightMost = rightMost->_right;}
_node = rightMost;
}else{
Node* cur = _node;
Node* parent = cur->_parent;
while(parent && cur == parent->_left){ cur = parent; parent = cur->_parent;}
_node = parent;
}
return *this;
}
- end() 迭代器指向 nullptr
- 当对 end() 执行–操作时,应该得到最后一个有效节点(最右节点)
- 所以需要记录_root 来找到最右节点
3.4 RBTree 中的迭代器接口
template<class K,class T,class KeyOfT>
class RBTree{
public:
typedef RBTreeIterator<T, T&, T*> Iterator;
typedef RBTreeIterator<T,const T&,const T*> ConstIterator;
Iterator Begin(){
Node* leftMost = _root;
while(leftMost && leftMost->_left){ leftMost = leftMost->_left;}
return Iterator(leftMost, _root);
}
Iterator End(){return Iterator(nullptr, _root);}
ConstIterator Begin()const{
Node* leftMost = _root;
while(leftMost && leftMost->_left){ leftMost = leftMost->_left;}
return ConstIterator(leftMost, _root);
}
ConstIterator End()const{return ConstIterator(nullptr, _root);}
};
四。map 和 set 对迭代器的封装
4.1 set 的迭代器封装
namespace bit {
template<class K>
class set{
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();}
private:
RBTree<K,const K, SetKeyOfT> _t;
};
}
4.2 map 的迭代器封装
namespace bit {
template<class K,class V>
class map{
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();}
private:
RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};
}
4.3 测试迭代器
void test_set(){
bit::set<int> s;
int a[]={4,2,6,1,3,5,15,7,16,14};
for(auto e : a){ s.insert(e); }
for(auto e : s){ cout << e <<" ";}
cout << endl;
}
五。map 的 operator[] 实现
5.1 operator[] 的原理
map 的 operator[] 是一个非常方便的特性:
- 如果 key 存在,返回对应的 value 引用
- 如果 key 不存在,插入一个默认 value 并返回其引用
5.2 实现步骤
- 修改 RBTree 的 Insert,返回
pair<Iterator, bool>
- 在 map 的 operator[] 中调用 insert
- 根据返回的迭代器访问 second
5.3 代码实现
namespace bit {
template<class K,class V>
class map{
public:
V& operator[](const K& key){
pair<iterator,bool> ret = insert(make_pair(key,V()));
return ret.first->second;
}
};
}
5.4 测试 operator[]
void test_map(){
bit::map<string, string> dict;
dict.insert({"sort","排序"});
dict.insert({"left","左边"});
dict.insert({"right","右边"});
dict["left"]="左边,剩余";
dict["insert"]="插入";
dict["string"];
for(auto& kv : dict){
kv.second +='x';
cout << kv.first <<":"<< kv.second << endl;
}
}
insert:插入 x left:左边,剩余 x right:右边 x sort:排序 x string:x
六。const 迭代器和 const 正确性
6.1 const 迭代器的使用场景
void Print(const bit::set<int>& s){
bit::set<int>::const_iterator it = s.begin();
while(it != s.end()){
cout <<*it <<" ";
++it;
}
cout << endl;
}
6.2 为什么要支持 const 迭代器?
- const 对象只能用 const 迭代器:保证 const 对象的内容不会被修改
- 语义清晰:表明这段代码只读不写
- 编译器检查:帮助捕获意外的修改
6.3 set 中 key 不可修改的保证
template<class K>
class set{
private:
RBTree<K,const K, SetKeyOfT> _t;
};
这样,迭代器解引用得到的是 const K&,无法修改。
6.4 map 中 key 不可修改,value 可修改
template<class K,class V>
class map{
private:
RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};
- 迭代器解引用得到
pair<const K, V>&
first 是 const,不能修改
second 不是 const,可以修改
七。完整代码实现
7.1 RBTree.h(完整版)
#pragma once
#include <utility>
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),_col(RED){}
};
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){}
Self& operator++(){
if(_node ==nullptr) return this;
if(_node->_right){
Node* leftMost = _node->_right;
while(leftMost->_left){ leftMost = leftMost->_left;}
_node = leftMost;
}else{
Node* cur = _node;
Node* parent = cur->_parent;
while(parent && cur == parent->_right){ cur = parent; parent = cur->_parent;}
_node = parent;
}
return *this;
}
Self& operator--(){
if(_node ==nullptr){
Node* rightMost = _root;
while(rightMost && rightMost->_right){ rightMost = rightMost->_right;}
_node = rightMost;
}else if(_node->_left){
Node* rightMost = _node->_left;
while(rightMost->_right){ rightMost = rightMost->_right;}
_node = rightMost;
}else{
Node* cur = _node;
Node* parent = cur->_parent;
while(parent && cur == parent->_left){ cur = parent; parent = cur->_parent;}
_node = parent;
}
return *this;
}
Ref operator*(){return _node->_data;}
Ptr operator->(){return&_node->_data;}
bool operator!=(const Self& s)const{return _node != s._node;}
bool operator==(const Self& s)const{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*> ConstIterator;
RBTree()=default;
~RBTree(){Destroy(_root); _root =nullptr;}
Iterator Begin(){
Node* leftMost = _root;
while(leftMost && leftMost->_left){ leftMost = leftMost->_left;}
return Iterator(leftMost, _root);
}
Iterator End(){return Iterator(nullptr, _root);}
ConstIterator Begin()const{
Node* leftMost = _root;
while(leftMost && leftMost->_left){ leftMost = leftMost->_left;}
return ConstIterator(leftMost, _root);
}
ConstIterator End()const{return ConstIterator(nullptr, _root);}
pair<Iterator,bool>Insert(const T& data){
if(_root ==nullptr){
_root =new Node(data);
_root->_col = BLACK;
return make_pair(Iterator(_root, _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)){ parent = cur; cur = cur->_left;}
else{return make_pair(Iterator(cur, _root),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;
while(parent && parent->_col == RED){
Node* grandfather = parent->_parent;
if(parent == grandfather->_left){
Node* uncle = grandfather->_right;
if(uncle && uncle->_col == RED){
parent->_col = 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 = 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 make_pair(Iterator(newnode, _root),true);
}
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, _root);
}
return End();
}
private:
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 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(parentParent ==nullptr){ _root = subL; subL->_parent =nullptr;}
else{
if(parent == parentParent->_left) parentParent->_left = subL;
else parentParent->_right = subL;
subL->_parent = parentParent;
}
}
void Destroy(Node* root){
if(root ==nullptr) return;
Destroy(root->_left);
Destroy(root->_right);
delete root;
}
Node* _root =nullptr;
};
7.2 Myset.h(完整版)
#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;
};
void Print(const set<int>& s){
set<int>::const_iterator it = s.begin();
while(it != s.end()){
cout <<*it <<" ";
++it;
}
cout << endl;
}
void test_set(){
set<int> s;
int a[]={4,2,6,1,3,5,15,7,16,14};
for(auto e : a){ s.insert(e); }
for(auto e : s){ cout << e <<" ";}
cout << endl;
Print(s);
}
}
7.3 Mymap.h(完整版)
#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 = insert(make_pair(key,V()));
return ret.first->second;
}
private:
RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};
void test_map(){
map<string, string> dict;
dict.insert({"sort","排序"});
dict.insert({"left","左边"});
dict.insert({"right","右边"});
dict["left"]="左边,剩余";
dict["insert"]="插入";
dict["string"];
map<string, string>::iterator it = dict.begin();
while(it != dict.end()){
it->second +='x';
cout << it->first <<":"<< it->second << endl;
++it;
}
cout << endl;
}
}
八。总结与思考
8.1 设计亮点
- 泛型设计:一棵红黑树通过模板参数适配 set 和 map
- 仿函数 KeyOfT:解决了不同类型中取 key 的问题
- 迭代器抽象:封装了中序遍历的逻辑
- const 正确性:通过模板参数 Ref 和 Ptr 支持 const 迭代器
8.2 与 STL 源码的对比
| 特性 | 我们的实现 | SGI-STL |
|---|
| 节点结构 | RBTreeNode | __rb_tree_node |
| 迭代器 | RBTreeIterator | __rb_tree_iterator |
| KeyOfT | 模板参数 | identity/select1st |
| 迭代器– | 支持 | 支持 |
| 哨兵节点 | 无 | header 节点 |
8.3 常见问题
- set 中的元素本身就是 key,修改 key 会破坏 BST 的有序性
- 实现中通过
RBTree<K, const K, SetKeyOfT> 保证存储的是 const K
Q2:为什么 map 的迭代器能修改 value 但不能修改 key?
- map 存储的是
pair<const K, V>,key 是 const 的,value 不是 const
- 修改 key 会破坏有序性,修改 value 不会影响树的组织结构
Q3:为什么要实现 const_iterator?
- const 对象只能用 const_iterator 遍历
- 明确表达只读语义,帮助编译器检查
- 为了处理
--end() 的情况,需要找到最右节点
- 没有_root 无法从 nullptr 找到最右节点
相关免费在线工具
- 加密/解密文本
使用加密算法(如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