封装红黑树实现map和set —— 从源码到模拟实现

一.、从 STL 源码看 map 和 set 的复用结构

1.1 为什么 map 和 set 可以共用一颗红黑树？

在 SGI-STL 中，set 和 map 并没有各自实现一棵红黑树，而是都依赖于一个通用的 rb_tree 模板类。这个类的设计非常巧妙------它通过模板参数来决定树节点中存储的数据类型，从而实现 "数据类型的泛化"。

我们来看源码中的关键片段（来自 stl_set.h 和 stl_map.h）：

cpp 复制代码

// 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;   // 底层红黑树
};

// 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;
};

注意rb_tree的模板参数：

cpp 复制代码

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

Key：键的类型，用于查找和删除。
Value ：节点中实际存储的数据类型。对于 set，它就是 Key；对于 map，它是 pair<const Key, T>。
KeyOfValue：一个仿函数，用于从 Value 中提取 Key，以便在比较时只比较键。

** rb_tree**不关心 Value 到底是什么，它只负责按 Key 排序并维护平衡。这样，set 和 map 就能复用同一份红黑树代码了。

1.2 为什么还要单独传递 Key 参数？

你可能发现 set 的 Key 和 Value 都是同一个类型，map 的 Key 和 Value 的 first 类型也相同，那为什么 rb_tree 还需要第一个参数 Key 呢？

这是因为 rb_tree 的 find、erase 等成员函数的参数类型是 key_type（即第一个模板参数），而不是 value_type。例如：

cpp 复制代码

iterator find(const key_type& x);
size_type erase(const key_type& x);

对于 set，查找时传入一个 Key 对象；对于 map，查找时也是传入 Key 对象（而不是整个 pair）。所以必须把 Key 单独拎出来。

1.3 红黑树节点的结构

SGI 中节点定义如下：

cpp 复制代码

struct __rb_tree_node_base {
    typedef __rb_tree_color_type color_type;
    color_type color;
    base_ptr parent;
    base_ptr left;
    base_ptr right;
};

template <class Value>
struct __rb_tree_node : public __rb_tree_node_base {
    Value value_field;   // 实际存储的数据
};

节点继承自基类，基类只包含颜色和三个指针，派生类多一个 value_field 。这样做的目的是为了节省内存？其实更多是为了分离类型无关的操作。不过我们模拟实现时，直接用单个**RBTreeNode**模板即可。

二、开始模拟实现 ------ 从红黑树节点开始

我们先从自己实现的代码出发，对照上面源码的思路。

2.1 颜色枚举和节点结构

这里我们和源码保持一致：节点持有 _data（即源码中的 value_field），并且带有三叉链（左、右、父），方便后续旋转和迭代器遍历。

三、红黑树的框架 ------ 引入 KeyOfT 仿函数

因为我们的红黑树现在并不知道 T 是什么类型，如果 T 是 pair<K,V>，我们比较时不能直接用 < 比较整个 pair（那样会同时比较 first 和 second），而应该只比较 first（即 key）。

为此，我们在上层（map 和 set）定义各自的仿函数，传递给红黑树，红黑树通过该仿函数来提取 key。

先看红黑树的模板声明：

cpp 复制代码

template<class K, class T, class KeyOfT>
class RBTree {
    // K 是键类型，T 是节点存储的数据类型，KeyOfT 是提取 key 的仿函数
    // ...
};

在插入时，我们使用**KeyOfT**对象来比较：

cpp 复制代码

KeyOfT kot;
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 false;
    }
}

这样，无论 T 是 K 还是**pair<K,V>**，我们都能够统一按 key 进行比较。

四、Insert 实现 ------ 按 key 插入，并进行平衡调整

4.1 基本插入逻辑

我们的 **Insert**函数返回 pair<Iterator, bool>，这样既能获得新插入节点的迭代器，又能知道是否插入成功。

4.2 平衡调整（红黑树核心）

调整的逻辑和标准红黑树一致（上篇文章我们详细讲过），分为几种情况（设当前节点为 cur，父节点为 parent，祖父为 grandfather，叔叔为 uncle）：

情况1：父节点存在且为黑 → 无需调整，直接结束。
情况2：父节点为红，叔叔存在且为红 → 将父和叔变黑，祖父变红，然后 cur 往上跳到祖父，继续检查。
情况3：父为红，叔叔不存在或为黑，且 cur 与 parent 在同侧（如都是左孩子） → 单旋 + 变色。
情况4：父为红，叔叔不存在或为黑，且 cur 与 parent 在异侧（如 cur 是右孩子，parent 是左孩子） → 双旋 + 变色。

我们的代码中，旋转函数 RotateL 和 RotateR 需要正确更新父指针。

在课件中，插入调整的代码如下：

最后返回迭代器，注意**Iterator**构造时需要传入当前节点和根节点（后面会讲为什么需要根）。

五、迭代器 ------ 中序遍历的封装

红黑树的迭代器本质是一个"节点指针的封装"，重载 ++ 和 -- 实现中序遍历的向前/向后移动。

5.1 迭代器类模板

我们模仿源码，设计一个模板类 RBTreeIterator，它有三个模板参数：T（数据类型）、Ref（引用类型）、Ptr（指针类型）。这样我们可以同时支持普通迭代器和 const 迭代器。

重点解析：

为什么还要存 _root？ 这是为了 --end() 操作。当我们对 end()（_node == nullptr）执行 -- 时，需要找到整棵树最右节点，而只有根节点指针才能遍历到最右。SGI 源码用了一个哨兵头节点，我们选择保存根节点来简化。
operator++ 逻辑：
- 若有右孩子，则下一个是右子树的最左节点；
- 若无右孩子，则沿父指针向上，直到当前节点是父节点的左孩子，那么父节点就是下一个；若一直找到根都未满足（即当前节点是整棵树的最右节点），则最终 _node 变为 nullptr，表示 end()。
operator-- 逻辑：完全对称，若有左孩子，则前一个是左子树的最右节点；若无左孩子，则向上找"当前节点是父节点右孩子"的祖先。

5.2 在 RBTree 中定义迭代器类型

Begin() 返回中序第一个节点（最左节点），End() 返回空指针。

六、封装 set ------ 去重且只存 key

6.1 SetKeyOfT 仿函数

在 set 中，T 就是 K，提取 key 直接返回自身即可。

特别注意 ：我们将 RBTree 的第二个模板参数（即节点存储的数据类型）设为 const K，这样通过迭代器解引用得到的是 const K&，用户无法修改 key，保证了 set 的元素不可变。

6.2 测试 set

七、封装 map ------ key-value 映射

7.1 MapKeyOfT 仿函数

在 map 中，T 是 pair<const K, V>，提取 key 需要返回 kv.first。

operator[] 实现 ：先调用 insert，如果 key 已存在，返回已有元素的迭代器；否则插入默认值（V()），然后返回对应 value 的引用。这样就能支持 dict["insert"] = "插入" 这样的语法。

7.2 测试 map

八、常见问题与注意事项

红黑树旋转函数：我们在上篇文章里具体讲过
迭代器 -- 处理 end() ：我们的迭代器保存了 _root，所以在 -- 时能特殊处理。如果 _node == nullptr，我们就让它指向最右节点。
set 不允许修改 key ：我们在实例化 RBTree 时，将 T 设为 const K，所以任何通过迭代器解引用得到的都是 const K&，无法修改。
map 的 key 不允许修改，value 可以 ：我们用 pair<const K, V> 作为 T，这样 first 是 const，second 是可修改的。
仿函数的传递 ：KeyOfT 必须在 RBTree 内部使用，并且每个比较操作都要调用 kot(data) 来提取 key。