红黑树进阶：从零实现 STL 容器
在 SGI-STL 的实现中，map 和 set 的底层核心是一棵红黑树（rb_tree）。这种设计体现了 C++ 泛型编程的强大之处：通过同一棵红黑树，既可以实现 key-only 的 set，也可以实现 key/value 的 map。
一、设计思想与框架
1.1 核心设计
红黑树的节点中存储什么类型，由模板参数决定。对于 set，传给 rb_tree 的是 Key 类型；对于 map，则是 pair<const Key, T> 类型。
这里有个关键点：为什么 rb_tree 需要两个模板参数 Key 和 Value？
- Value：决定了节点中存储的数据类型（set 是 Key，map 是 pair）。
- Key：用于查找和删除操作的参数类型（find/erase 的参数类型）。
对于 set 而言，Key 和 Value 是相同的；对于 map 而言，Key 是键的类型，Value 是 pair 类型。这种设计使得一棵红黑树既能处理 set 场景，也能处理 map 场景。
1.2 源码框架
// stl_set.h
template<class Key, class Compare = less<Key>, class Alloc = alloc>
class set {
public:
typedef Key key_type;
typedef Key value_type; // set 的 value_type 就是 Key
private:
typedef rb_tree<key_type, value_type, identity<value_type>, key_compare, Alloc> rep_type;
rep_type t; // 红黑树，存储的是 Key 类型
};
// 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; // map 的 value_type 是 pair
private:
typedef rb_tree<key_type, value_type, select1st<value_type>, key_compare, Alloc> rep_type;
rep_type t; // 红黑树，存储的是 pair<const Key, T> 类型
};
二、模拟实现 map 和 set
2.1 红黑树节点定义
// RBTree.h
#pragma once
enum Colour { RED, BLACK };
template<class T>
struct RBTreeNode {
T _data; // 节点存储的数据（泛型 T）
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 进行比较。
// Mymap.h - map 中定义的 KeyOfT 仿函数
namespace bit {
template<class K, class V>
class map {
struct MapKeyOfT {
const K& operator()(const pair<K, V>& kv) {
return kv.first; // 返回 pair 的 first（键）
}
};
// ... 其他代码
};
}
// Myset.h - set 中定义的 KeyOfT 仿函数
namespace bit {
template<class K>
class set {
struct SetKeyOfT {
const K& operator()(const K& key) {
return key; // 直接返回 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 封装
// Mymap.h
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;
};
}
// Myset.h
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-> 以及比较操作。
template<class T, class Ref, class Ptr>
struct RBTreeIterator {
typedef RBTreeNode<T> Node;
typedef RBTreeIterator<T, Ref, Ptr> Self;
Node* _node;
Node* _root; // 根节点（用于 --end() 特殊处理）
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;
}
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() 的特殊处理：end() 指向 nullptr，执行 -- 时应得到最后一个有效节点（最右节点），所以需要记录 _root。
3.3 接口封装
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 类型重命名与测试
// Myset.h
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;
};
}
// 测试
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 << " "; // 输出：1 2 3 4 5 6 7 14 15 16
cout << endl;
}
五、map 的 operator[] 实现
5.1 原理与步骤
map 的 operator[] 非常方便：如果 key 存在，返回对应的 value 引用；如果不存在，插入一个默认 value 并返回其引用。
实现步骤：
- 修改 RBTree 的 Insert，返回
pair<Iterator, bool>。 - 在 map 的 operator[] 中调用 insert。
- 根据返回的迭代器访问 second。
5.2 代码实现
// Mymap.h
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;
}
private:
RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};
}
5.3 测试
void test_map() {
bit::map<string, string> dict;
dict.insert({"sort","排序"});
dict["left"]="左边，剩余";
dict["insert"]="插入";
dict["string"]; // 插入新 key，value 为默认值
for(auto& kv : dict) {
kv.second += 'x'; // 正确：value 可以修改
cout << kv.first << ":" << kv.second << endl;
}
}
六、const 迭代器和 const 正确性
6.1 使用场景与必要性
接收 const 引用的函数只能用 const_iterator，因为 s 是 const 的。这保证了 const 对象的内容不会被修改，语义清晰且能帮助编译器检查。
6.2 保证 Key 不可修改
在 set 中，key 是不可修改的。实现方式是将 RBTree 第二个模板参数传入 const K。
// Myset.h
template<class K>
class set {
private:
RBTree<K, const K, SetKeyOfT> _t; // 存储 const K
};
这样，迭代器解引用得到的是 const K&，无法修改。
6.3 Map 中的 Key 与 Value
// Mymap.h
template<class K, class V>
class map {
private:
RBTree<K, pair<const K, V>, MapKeyOfT> _t; // key 是 const 的，value 是可修改的
};
迭代器解引用得到 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;
};
}
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;
};
}
八、总结与思考
8.1 设计亮点
- 泛型设计：一棵红黑树通过模板参数适配 set 和 map。
- 仿函数 KeyOfT：解决了不同类型中取 key 的问题。
- 迭代器抽象：封装了中序遍历的逻辑。
- const 正确性：通过模板参数 Ref 和 Ptr 支持 const 迭代器。
8.2 常见问题
Q1：为什么 set 的迭代器不能修改 key？
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 遍历，明确表达只读语义，帮助编译器检查。
Q4：迭代器 -- 操作中为什么要保存_root？
为了处理 --end() 的情况，需要找到最右节点。没有_root 无法从 nullptr 找到最右节点。