从零手写 STL Set/Map：基于红黑树的泛型设计与实现

通过框架分析可以看到，rb_tree 使用了巧妙的泛型思想。它不直接写死存储类型，而是由第二个模板参数 Value 决定 _rb_tree_node 中存储的数据类型。

• set 实例化时，第二个模板参数给的是 Key。
• map 实例化时，第二个模板参数给的是 pair<const Key, T>。

这样一颗红黑树既可以实现 key 搜索场景，也可以实现 key/value 搜索场景的 map。至于为什么还要传第一个模板参数 Key？这是因为对于 map 和 set，find/erase 时的函数参数都是 Key，所以第一个模板参数是传给这些函数做形参类型的。对于 set 而言两个参数是一样的，但对于 map 而言就完全不一样了，map insert 的是 pair 对象，但是 find 和 erase 的 Key 对象。

二、核心设计思路：红黑树的泛型复用

STL 中 map 和 set 复用同一颗红黑树的核心是泛型编程 + 仿函数提取 key，解决了'一颗树适配两种数据场景'的问题。

2.1 红黑树的模板参数设计

红黑树需要支持存储两种数据类型：

• set 场景：存储单个 key（如 int、string）；
• map 场景：存储 pair<const Key, Value>（key 不可修改）。

因此红黑树的模板参数需抽象为 3 个：

template<class K, class T, class KeyOfT>
class RBTree {
    // K：find/erase 时的 key 类型（统一接口参数）
    // T：红黑树节点存储的实际数据类型（set 为 K，map 为 pair<const K, V>）
    // KeyOfT：仿函数，从 T 中提取 K（解决 T 类型不统一的比较问题）
};

2.2 仿函数 KeyOfT：统一 key 提取逻辑

由于 T 的类型不固定（K 或 pair），红黑树插入 / 查找时无法直接获取 key，需通过仿函数 KeyOfT 统一提取，由 map 和 set 分别实现适配：

• set 的仿函数：直接返回 key（T=K）；
• map 的仿函数：返回 pair 的 first 成员（T=pair<const K, V>）。

2.3 核心约束：key 不可修改

• set 的 key 是唯一标识，需禁止修改：红黑树存储 const K；
• map 的 key 是索引，需禁止修改：pair 的 first 设为 const K，value 可正常修改。

三、基础组件实现：红黑树与仿函数

3.1 红黑树节点结构

节点存储模板类型 T，包含左右子指针、父指针和颜色标记：

#pragma once
#include <iostream>
#include <assert.h>
using namespace std;

enum Colour { Red, Black };

template<class T>
struct RBTreeNode {
    T _data;
    RBTreeNode<T>* _parent;
    RBTreeNode<T>* _left;
    RBTreeNode<T>* _right;
    Colour _col;

    RBTreeNode(const T& data) : _parent(nullptr), _left(nullptr), _right(nullptr), _data(data), _col(Red) {}
};

3.2 仿函数实现（map/set 层）

3.2.1 set 的仿函数：直接返回 key

#pragma once
#include "RBTree.h"
namespace Lotso {
template<class K>
class set {
    struct SetKeyofT {
        const K& operator()(const K& key) { return key; }
    };
private:
    RBTree<K, const K, SetKeyofT> _t;
};
}

3.2.2 map 的仿函数：提取 pair 的 first

#pragma once
#include "RBTree.h"
namespace Lotso {
template<class K, class V>
class map {
    struct MapKeyofT {
        const K& operator()(const pair<K, V>& kv) { return kv.first; }
    };
private:
    RBTree<K, pair<const K, V>, MapKeyofT> _t;
};
}

3.3 红黑树核心接口（附迭代器）

重点实现 Insert（返回 pair<Iterator, bool>，支持 map 的 []）和 Find，平衡维护逻辑与基础红黑树一致。这里展示一下 operator++ 的实现，-- 的话就不展示了，要实现的话还需要额外带一个 _root。

#pragma once
#include <iostream>
#include <assert.h>
using namespace std;

enum Colour { Red, Black };

template<class T>
struct RBTreeNode {
    T _data;
    RBTreeNode<T>* _parent;
    RBTreeNode<T>* _left;
    RBTreeNode<T>* _right;
    Colour _col;
    RBTreeNode(const T& data) : _parent(nullptr), _left(nullptr), _right(nullptr), _data(data), _col(Red) {}
};

template<class T, class Ref, class Ptr>
struct RBTreeIterator {
    typedef RBTreeNode<T> Node;
    typedef RBTreeIterator<T, Ref, Ptr> Self;
    Node* _node;

    RBTreeIterator(Node* node) : _node(node) {}

    Self& operator++() {
        if (_node->_right) {
            Node* minRight = _node->_right;
            while (minRight->_left) minRight = minRight->_left;
            _node = minRight;
        } else {
            Node* cur = _node;
            Node* parent = cur->_parent;
            while (parent && cur == parent->_right) {
                cur = parent;
                parent = parent->_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; }
};

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() { Destory(_root); _root = nullptr; }

    void Destory(Node* root) {
        if (root == nullptr) return;
        Destory(root->_left);
        Destory(root->_right);
        delete root;
    }

    Iterator Begin() {
        Node* minLeft = _root;
        while (minLeft && minLeft->_left) minLeft = minLeft->_left;
        return Iterator(minLeft);
    }

    Iterator End() { return Iterator(nullptr); }

    ConstIterator Begin() const {
        Node* minLeft = _root;
        while (minLeft && minLeft->_left) minLeft = minLeft->_left;
        return ConstIterator(minLeft);
    }

    ConstIterator End() const { return ConstIterator(nullptr); }

    pair<Iterator, bool> Insert(const T& data) {
        if (_root == nullptr) {
            _root = new Node(data);
            _root->_col = Black;
            return { Iterator(_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(data) < kot(cur->_data)) {
                parent = cur;
                cur = cur->_left;
            } else {
                return { Iterator(cur), 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* grandparent = parent->_parent;
            if (grandparent->_left == parent) {
                Node* uncle = grandparent->_right;
                if (uncle && uncle->_col == Red) {
                    parent->_col = Black;
                    uncle->_col = Black;
                    grandparent->_col = Red;
                    cur = grandparent;
                    parent = cur->_parent;
                } else {
                    if (cur == parent->_left) {
                        RotateR(grandparent);
                        parent->_col = Black;
                        grandparent->_col = Red;
                    } else {
                        RotateL(parent);
                        RotateR(grandparent);
                        cur->_col = Black;
                        grandparent->_col = Red;
                    }
                    break;
                }
            } else {
                Node* uncle = grandparent->_left;
                if (uncle && uncle->_col == Red) {
                    uncle->_col = Black;
                    parent->_col = Black;
                    grandparent->_col = Red;
                    cur = grandparent;
                    parent = cur->_parent;
                } else {
                    if (parent->_right == cur) {
                        RotateL(grandparent);
                        parent->_col = Black;
                        grandparent->_col = Red;
                    } else {
                        RotateR(parent);
                        RotateL(grandparent);
                        cur->_col = Black;
                        grandparent->_col = Red;
                    }
                    break;
                }
            }
        }
        _root->_col = Black;
        return { Iterator(newnode), 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);
            }
        }
        return End();
    }

private:
    void RotateR(Node* parent) {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        parent->_left = subLR;
        if (subLR) subLR->_parent = parent;
        Node* grandparent = parent->_parent;
        subL->_right = parent;
        parent->_parent = subL;
        if (parent == _root) {
            _root = subL;
            subL->_parent = nullptr;
        } else {
            if (grandparent->_left == parent) grandparent->_left = subL;
            else grandparent->_right = subL;
            subL->_parent = grandparent;
        }
    }

    void RotateL(Node* parent) {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        parent->_right = subRL;
        if (subRL) subRL->_parent = parent;
        Node* grandparent = parent->_parent;
        subR->_left = parent;
        parent->_parent = subR;
        if (_root == parent) {
            _root = subR;
            subR->_parent = nullptr;
        } else {
            if (grandparent->_left == parent) grandparent->_left = subR;
            else grandparent->_right = subR;
            subR->_parent = grandparent;
        }
    }

private:
    Node* _root = nullptr;
};

3.4 迭代器实现思路分析

迭代器实现的大框架跟 list 的 iterator 思路是一致的，用一个类型封装结点的指针，再通过重载运算符实现，迭代器像指针一样访问的行为。

这里的难点是 operator++ 和 operator-- 的实现。之前使用部分，我们分析了，map 和 set 的迭代器走的是中序遍历，左子树 -> 根结点 -> 右子树，那么 begin() 会返回中序第一个结点的 iterator。

• 迭代器 ++ 的核心逻辑就是不看全局，只看局部，只考虑当前中序局部要访问的下一个结点。
• 如果 it 指向的结点的右子树不为空，代表当前结点已经访问完了，要访问下一个结点是右子树的中序第一个，一棵树中序第一个是最左结点，所以直接找右子树的最左结点即可。
• 如果 it 指向的结点的右子树空，代表当前结点已经访问完了且当前结点所在的子树也访问完了，要访问的下一个结点在当前结点的祖先里面，所以要沿着当前结点到根的祖先路径向上找。
• 如果当前结点是父亲的左，根据中序左子树 -> 根结点 -> 右子树，那么下一个访问的结点就是当前结点的父亲。
• 如果当前结点时父亲的右，根据中序左子树 -> 根结点 -> 右子树，当前结点所在的子树访问完了，当前结点所在父亲的子树也已经访问完了，那么下一个访问的需要继续往根的祖先中去找，直到找到孩子是父亲左的那个祖先就是中序要走的下一个结点。
• end() 如何表示呢？当 it 指向最后一个结点时，++it 时，一直往上找没有孩子是父亲左的那个祖先，这时父亲为空了，那么我们就把 it 中的结点指针置为 nullptr，我们用去充当 end。需要注意的是 STL 源码中，红黑树增加了一个哨兵位头结点做为 end()，这哨兵位头结点和根互为父亲，左指向最左结点，右指向最右结点。相比我们用 nullptr 作为 end()，差别不大，他能实现的，我们也能实现。只是 --end() 判断到结点是空，特殊处理一下，让迭代器结点指向最右结点。

set 的 iterator 也不支持修改，我们把 set 的第二个模板参数改成 const K 即可，RBTree<K, const K, SetKeyOfT> _t； map 的 iterator 不支持修改 key 但是可以修改 value，我们把 map 的第二个模板参数 pair 的第一个参数改成 const K 即可，RBTree<K, pair<const K, V>, MapKeyOfT> _t；

支持完整的迭代器还有很多细节需要修改，具体参考上面的代码。

四、mySet 与 myMap 完整实现

map 支持 [] 主要修改 insert 返回值支持，修改 RBTree 中的 insert 返回值为 pair<Iterator, bool> Insert(const T& data)。

4.1 mySet 实现

#pragma once
#include "RBTree.h"
namespace Lotso {
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;
};
}

4.2 myMap 实现

#pragma once
#include "RBTree.h"
namespace Lotso {
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) {
        auto [it, flag] = _t.Insert({ key, V() });
        return it->second;
    }

private:
    RBTree<K, pair<const K, V>, MapKeyofT> _t;
};
}

五、测试代码：验证功能

#define _CRT_SECURE_NO_WARNINGS 1
#include "RBTree.h"
#include "Map.h"
#include "Set.h"

template<class T>
void func(const Lotso::set<T>& s) {
    typename Lotso::set<T>::const_iterator it = s.begin();
    while (it != s.end()) {
        cout << *it << " ";
        ++it;
    }
    cout << endl;
}

void test_set() {
    Lotso::set<int> s;
    s.insert(1); s.insert(2); s.insert(1); s.insert(5);
    s.insert(0); s.insert(10); s.insert(8);
    Lotso::set<int>::iterator it = s.begin();
    while (it != s.end()) {
        cout << *it << " ";
        ++it;
    }
    cout << endl;
    func(s);
}

void test_map() {
    Lotso::map<string, string> dict;
    dict.insert({"sort", "排序"});
    dict.insert({"left", "左边"});
    dict.insert({"right", "右边"});
    dict["string"] = "字符串";
    dict["left"] = "左边 xxx";

    auto it = dict.begin();
    while (it != dict.end()) {
        it->second += 'x';
        cout << it->first << ":" << it->second << endl;
        ++it;
    }
    cout << endl;

    for (auto& [k, v] : dict) {
        cout << k << ":" << v << endl;
    }
    cout << endl;

    string arr[] = {"苹果", "西瓜", "苹果", "西瓜", "苹果", "苹果", "西瓜", "苹果", "香蕉", "苹果", "香蕉"};
    Lotso::map<string, int> countMap;
    for (auto& e : arr) {
        countMap[e]++;
    }
    for (auto& [k, v] : countMap) {
        cout << k << ":" << v << endl;
    }
    cout << endl;
}

int main() {
    cout << "测试 set:" << endl;
    test_set();
    cout << "------------------" << endl;
    cout << "测试 map:" << endl;
    test_map();
    return 0;
}

测试没有问题，可以正常使用。

结语

myMap 和 mySet 的封装核心是红黑树的泛型复用：通过模板参数 T 适配不同数据类型，用仿函数 KeyOfT 统一 key 提取逻辑，再通过迭代器封装实现双向遍历。其中，key 不可修改的约束、map 的 [] 运算符重载、迭代器的 ++/-- 实现，都是 STL 容器设计的经典细节。掌握这套封装逻辑，不仅能理解 STL 容器的底层实现，更能学会'泛型编程 + 接口抽象'的设计思想——用一套核心代码适配多种场景，兼顾灵活性与效率。