手撕 STL:一棵红黑树,凭什么同时撑起 set 和 map?

摘要:学红黑树时几乎人人都有个疑问------std::set 存 key,std::map 存 pair<key, value>,存的东西根本不一样,为什么它们底层是同一棵 红黑树?本文先从 SGI STL 源码切入,看 STL 官方是怎么设计的,再用一份可编译的 C++ 代码复刻它:讲透模板参数 K/V、KeyOfT 仿函数、以及迭代器 operator++ 的中序后继。读完你能自己写出一个"能当 set、能当 map、能迭代器遍历"的红黑树封装。
📖 前置阅读 :本文默认你已理解红黑树的插入、变色与旋转。如果还不熟,建议先看我的上一篇------ 《手撕红黑树:插入、变色与旋转全解析》本文只聚焦"如何用这棵红黑树封装出 set 和 map",不再重复红黑树本身的调整细节。

前言:先看 STL 源码是怎么写的

在动手之前,先翻一眼 SGI STL 的源码。set 和 map 的底层类型定义分别藏在 stl_set.h 和 stl_map.h 里:

cpp 复制代码

// stl_set.h ------ set 的底层
template <class Key, class Compare = less<Key>, class Alloc = alloc>
class set {
    // value 就是 key 本身,取键用 identity(原样返回)
    typedef rb_tree<key_type, value_type,
                    identity<value_type>, key_compare, Alloc> rep_type;
    rep_type t;   // 一棵红黑树
};

// stl_map.h ------ map 的底层
template <class Key, class T, class Compare = less<Key>, class Alloc = alloc>
class map {
    // value 是 pair,取键用 select1st(取 pair.first)
    typedef rb_tree<key_type, value_type,
                    select1st<value_type>, key_compare, Alloc> rep_type;
    rep_type t;   // 同样是一棵红黑树
};

两者共用的红黑树模板长这样:

cpp 复制代码

// stl_tree.h
template <class Key, class Value, class KeyOfValue,
          class Compare, class Alloc = alloc>
class rb_tree { ... };

再看那两个"取键"的仿函数,它们就是 STL 定义在 stl_function.h 里的普通函数对象:

cpp 复制代码

template <class T>
struct identity {                       // set 用它
    const T& operator()(const T& x) const { return x; }
};

template <class Pair>
struct select1st {                      // map 用它
    const typename Pair::first_type&
    operator()(const Pair& x) const { return x.first; }
};

一切都清楚了:set 和 map 用的是同一个 rb_tree,区别只在给它传的模板参数不同 ------尤其是那个 KeyOfValue(如何从节点数据里取出键)。

classDiagram class set~Key~ { <<对外容器>> +Tree _root 红黑树根节点 +Insert(val) 转调 _root.Insert +begin() / end() } class map~Key, Value~ { <<对外容器>> +Tree _root 红黑树根节点 +Insert(pair) 转调 _root.Insert +begin() / end() } class RBTree~K, V, KeyOfT~ { -Node* _root 树根 +Insert(V) 插入并变色/旋转 +Find(K) 按键查找 +begin() / end() 中序迭代器 } class KeyOfT_set { <> +operator()(Key x) 原样返回 x } class KeyOfT_map { <> +operator()(pair x) 返回 x.first } set ..> RBTree : 组合 RBTree map ..> RBTree : 组合 RBTree, KeyOfT> RBTree ..> KeyOfT_set : set 传入的第三个模板参数 RBTree ..> KeyOfT_map : map 传入的第三个模板参数

解读:

set / map 是对外的壳 ,本身不存数据;各自持有一个红黑树成员 Tree _root(即 STL 源码里的 rep_type t),Insert/begin 等操作全部转调给它。

RBTree 是相同的底座 :set 和 map 各自拥有一棵独立的 红黑树实例(并非共享同一个对象),它们相同的是这同一份类模板代码 ------差异只在传入的模板参数:set 传 <Key, Key, KeyOfT=identity>,map 传 <Key, pair<K,V>, KeyOfT=select1st>。

KeyOfT 是取键仿函数 ------set 版原样返回键(对应 STL 的 identity),map 版返回 pair.first(对应 STL 的 select1st),它就是屏蔽 set/map 差异的关键。

为了对齐 STL,我把那个取键仿函数统一命名为 KeyOfT(即 STL 的 KeyOfValue)。看完你会发现,我们手写的 myset::KeyOfT 就是 identity,mymap::KeyOfT 就是 select1st。

1. 红黑树节点:三叉链 + 一个颜色位

先看节点。和普通 BST 的区别只有两点:多了一个指向父亲的指针(三叉链),多了一个颜色。

cpp 复制代码

enum NodeColor { red, black };

template<class T>
struct RBTreeNode
{
    using Node = RBTreeNode<T>;
    T      _data;      // 节点存的数据(类型是模板 T,不写死 K 或 pair)
    Node*  _left;
    Node*  _right;
    Node*  _parent;    // 三叉链:向上找父亲、叔叔、爷爷都要用
    NodeColor _color;

    RBTreeNode(const T& data)
        : _data(data), _left(nullptr), _right(nullptr)
        , _parent(nullptr), _color(red)   // 新节点默认红色
    {}
};

两个要点先埋下:

_data 的类型是 T,不是 K,也不是 pair。 这是"对存储无知"的第一步------节点只知道自己存了个 T,不关心它是 key 还是 pair。这正对应 STL 里 rb_tree 的第二个模板参数 Value。
新节点默认是红色。 新增红色,最多破坏"父子不能同红",影响面小;新增黑色则必然破坏"每条路径黑高相等",几乎一定要调整。

为什么要 _parent?红黑树插入后要向上回溯,靠"叔叔节点"的颜色决定变色还是旋转。没有父指针,每次都得从根往下找,得不偿失。

2. 关键抽象:三个模板参数各管什么

我们手写的红黑树类,模板声明是理解全文的钥匙(对应 STL 的 rb_tree,只是精简了 Compare/Alloc):

cpp 复制代码

template<class K, class V, class KeyOfT>
class RBTree {
public:
    using Node     = RBTreeNode<V>;
    using Iterator = RBTreeIterator<V, V&, V*>;
    // ...
};

三个参数各司其职,别混淆:

参数	含义	set 传入	map 传入	对应 STL
`K`	键类型,用于 `Find` 这类按 key 查找的接口	`Key`	`Key`	`Key`
`V`	节点实际存储的类型,也是 `value_type`	`Key`	`pair<Key,Value>`	`Value`
`KeyOfT`	仿函数类型,负责"从 V 里抽出 K"	`identity`	`select1st`	`KeyOfValue`

对 set:K == V,节点存的就是 key 本身。
对 map:K 是键,V 是 pair<K, V>,节点存整个键值对。

问题随之而来:红黑树内部比较大小时需要 key ,可它手里只有 V(可能是 pair)。怎么办?------这就是第三个参数 KeyOfT 的意义。

3. `KeyOfT` 仿函数

先看 set 的封装,它的 KeyOfT 等价于 STL 的 identity:

cpp 复制代码

template<class Key>
class myset {
public:
    struct KeyOfT {                                   // ≙ STL identity
        const Key& operator()(const Key& obj) { return obj; }   // 我自己就是 key
    };
    using value_type = Key;
    using Tree = RBTree<Key, value_type, KeyOfT>;     // K == V == Key
private:
    Tree _root;
};

再看 map 的封装,它的 KeyOfT 等价于 STL 的 select1st:

cpp 复制代码

template<class Key, class Value>
class mymap {
public:
    using value_type = std::pair<Key, Value>;
    struct KeyOfT {                                   // ≙ STL select1st
        const Key& operator()(const value_type& obj) { return obj.first; }  // 取 pair.first
    };
    using Tree = RBTree<Key, value_type, KeyOfT>;     // V 是 pair
private:
    Tree _root;
};

两个 KeyOfT 的 operator() 干的是同一件事------提取出Value的key :set 版原样返回,map 版返回 .first。

于是红黑树内部只要这样写,就对 set/map 完全通用:

cpp 复制代码

KeyOfT t;                       // 实例化仿函数
if (t(kv) < t(cur->_data))      { /* 往左走 */ }
else if (t(kv) > t(cur->_data)) { /* 往右走 */ }

t(...) 这个括号,就是 STL 屏蔽 set/map 差异的地方。红黑树永远只对 t(数据) 做比较,它根本不需要知道 _data 到底是 key 还是 pair。

flowchart LR A["set 插入 key (裸数字)"] --> T B["map 插入 pair (key, value 盒子)"] --> T T["KeyOfT 转接头 从数据中取出 key"] -->|"set: 原样返回 key"| C T -->|"map: 返回 pair.first"| C C["红黑树 只认 key 做比较 (全程不变)"]

只需实现不同类型的仿函数并传入模版参数，仿函数会自动取出key这就是把"如何取键"这个策略当参数传进来的泛型思想。

4. `Insert`:BST 插入 + 复用红黑树的修复

插入分两步:先按 BST 规则找位置插入,再向上修复红黑树性质。

4.1 先当普通 BST 插入

cpp 复制代码

void Insert(const V& kv) {
    if (_root == nullptr) {           // 空树:根必须是黑
        _root = new Node(kv);
        _root->_color = black;
        return;
    }

    KeyOfT t;
    Node* cur = _root;
    Node* parent = nullptr;
    while (cur) {
        if (t(kv) < t(cur->_data))      { parent = cur; cur = cur->_left;  }
        else if (t(kv) > t(cur->_data)) { parent = cur; cur = cur->_right; }
        else return;                    // key 已存在,去重(set/map 都不允许重复键)
    }

    Node* newnode = new Node(kv);       // 默认红色
    if (t(newnode->_data) < t(parent->_data)) parent->_left  = newnode;
    else                                       parent->_right = newnode;
    newnode->_parent = parent;
    // ... 接下来向上修复
}

比较全程用 t(...),这就是上一节的通用性在起作用。

4.2 向上修复:交给红黑树自己

插入后若出现"父子连续红",就要靠变色 / 旋转 修复红黑树性质。这部分逻辑(叔叔红变色上传、叔叔黑旋转 + 变色)与本文主题------如何封装 set/map ------是正交的:它属于红黑树自身的实现,set 和 map 完全不关心。

📌 红黑树的插入调整、变色与旋转细节,我已在上一篇写透,这里不再展开: 《手撕红黑树:插入、变色与旋转全解析》 👉 blog.csdn.net/jsjs1346/ar...

对本文而言只需记住一个结论:修复过程只动颜色和指针结构,从不关心 _data 里存的是 key 还是 pair。所以同一份修复代码,对 set 和 map 天然通用------这正是"封装"能成立的前提。

5. 迭代器:`operator++` 如何做中序后继

set/map 遍历出来是有序的,靠迭代器按中序走。中序递归很简单,但迭代器要"走一步停一步",所以必须把"找中序后继"拆成可迭代逻辑。

先看 begin() / end():

cpp 复制代码

Iterator begin() {              // 中序第一个 = 最左节点
    Node* cur = _root;
    while (cur && cur->_left) cur = cur->_left;
    return Iterator(cur);
}
Iterator end() { return Iterator(nullptr); }   // 用 nullptr 作尾后

核心是 operator++,找当前节点的中序后继,分两种情况:

cpp 复制代码

Self& operator++() {
    if (_node->_right != nullptr) {
        // 情况一:有右子树 → 后继是右子树的最左节点
        Node* cur = _node->_right;
        while (cur->_left) cur = cur->_left;
        _node = cur;
    } else {
        // 情况二:没有右子树 → 沿父链上溯,
        //         直到"我是父亲的左孩子",那个父亲就是后继
        Node* cur = _node;
        Node* parent = _node->_parent;
        while (parent && cur == parent->_right) {
            cur = parent;
            parent = parent->_parent;
        }
        _node = parent;
    }
    return *this;
}

理解这两种情况是本节关键:

有右子树:中序里"根之后紧接着右子树",而右子树里最先访问的是它的最左节点。
没有右子树 :当前子树走完,回到祖先。若我是父亲的右孩子,说明父亲也访问过了 ,继续往上;直到"我是父亲的左孩子",那个父亲才是未访问的后继。走到根还没找到,parent 变 nullptr------正好等于 end(),遍历自然终止。

配合 main.cpp 的用法就一目了然:

cpp 复制代码

mymap<std::string, std::string> mp;
mp.Insert({ "apple",   "苹果" });
mp.Insert({ "banana",  "香蕉" });
mp.Insert({ "abandon", "放弃" });

auto it = mp.begin();
while (it != mp.end()) {           // nullptr 作尾后,循环自然结束
    std::cout << it->first << ":" << it->second << " ";
    ++it;
}
// 输出按 key 字典序:abandon:放弃 apple:苹果 banana:香蕉

迭代器模板 RBTreeIterator<T, Ref, Ptr> 带了 Ref(引用)和 Ptr(指针)两个参数,是为了一套代码同时生成 iterator 和 const_iterator:普通迭代器传 V&, V*,常量迭代器传 const V&, const V*。这正是 STL 的标准手法。

6. 易错点总结

混淆 K 和 V。 K 是"查找用的键",V 是"节点实际存的东西"。set 里两者相同,map 里 V 是 pair,别把 Find(const K&) 的参数写成 V。
忘了用 KeyOfT 取键,直接拿 _data 比较。 对 map 来说 _data 是 pair,pair < pair 会连 value 一起比,语义就错了。所有比较必须走 t(_data)。
end() 用什么。 这里用 nullptr 作尾后,好处是 operator++ 上溯到根之上恰好得到 nullptr,无需哨兵。代价是 end() 不能 --。工业级 STL 用**带哨兵头节点(header)**的方案支持 --end(),这是可优化点。
operator++ 里 parent 的初始化。 只有上溯分支才用得到,建议按分支就近声明,可读性更好。
新节点颜色。 一律给红色,别设黑,否则每次插入都破坏黑高。

7. 总结

回到开头:set 和 map 凭什么共用一棵红黑树?

答案就藏在 SGI STL 的三层抽象里,而我们完整复刻了它:

模板参数 Value(本文 V):让节点对"存 key 还是存 pair"无知;
仿函数 KeyOfValue(本文 KeyOfT = identity/select1st) :把"如何取键"外包出去,红黑树只认 t(data);
统一迭代器 :用中序 operator++ 提供有序遍历,set/map 共享。

这套"把变化点抽象成参数 "的思想是泛型编程的通用范式。理解它之后再看 unordered_map(换哈希桶)、priority_queue(换堆),你会发现套路是相通的。

手撕 STL:一棵红黑树,凭什么同时撑起 set 和 map?

前言:先看 STL 源码是怎么写的

1. 红黑树节点:三叉链 + 一个颜色位

2. 关键抽象:三个模板参数各管什么

3. KeyOfT 仿函数

4. Insert:BST 插入 + 复用红黑树的修复

4.1 先当普通 BST 插入

4.2 向上修复:交给红黑树自己

5. 迭代器:operator++ 如何做中序后继

6. 易错点总结

7. 总结

3. `KeyOfT` 仿函数

4. `Insert`:BST 插入 + 复用红黑树的修复

5. 迭代器:`operator++` 如何做中序后继