C++ STL set/map 模拟实现

红黑树（rb_tree）泛型设计思想解析 通过对框架的分析，SGI STL 源码中 rb_tree 的泛型设计非常巧妙。它不直接写死'仅支持 key 搜索'或'仅支持 key/value 搜索'，而是通过第二个模板参数 Value 灵活控制：红黑树节点（__rb_tree_node）中实际存储的数据类型，由 Value 决定。 这样一颗红黑树，既能适配 set 的'纯 key 搜索场景'，也能适配 map 的'key/value 搜索场景'。

• set 的场景：当 set 实例化 rb_tree 时，给第二个模板参数传入纯 key 类型（如：set<int> 中，Value 就是 int）。红黑树节点直接存储 key，自然适配'仅按 key 搜索、不重复存储'的需求。
• map 的场景：当 map 实例化 rb_tree 时，给第二个模板参数传入键值对类型 pair<const Key, T>（如：map<int, string> 中，Value 是 pair<const int, string>）。红黑树节点存储完整的键值对，从而支持'按 key 关联 value'的搜索逻辑。

注意事项：关于 value_type 的特殊含义 源码里的模板参数常用 T 代表'节点存储的数据类型（即这里的 Value）'。而 rb_tree 内部定义的 value_type，并非我们日常说的'key/value 里的 value'，而是红黑树节点实际存储的数据的类型（对 set 是 key 类型，对 map 是 pair<const Key, T> 类型）。

2. 为什么需要两个模板参数（Key & Value）？

核心原因在于 find/erase 等操作的入参需求：

• 对 set 和 map 来说：find/erase 函数的参数是 Key 类型（按 key 查找、删除），而非完整的节点数据（Value 类型）。
• 对 set 而言：Key 和 Value 类型相同（节点存 key，操作也用 key），所以两个模板参数看似冗余，但是这样做主要是为了和 map 容器保持统一的接口。
• 对 map 而言：Key（操作入参类型）和 Value（节点存储的键值对类型）完全不同 —— map 插入的是 pair 对象，但查找、删除只用 Key。

因此：Key 模板参数的意义是为 find/erase 等函数提供形参类型，让红黑树能统一支撑 set（Key 与 Value 同类型）和 map（Key 与 Value 不同类型）的场景。

模拟实现

头文件

RBTree.h

#pragma once
#include <iostream>
#include <vector>
#include <time.h>
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) {}

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

    Self operator++() {
        if (_node->_right) {
            Node* min = _node->_right;
            while (min->_left) min = min->_left;
            _node = min;
        } else {
            Node* current = _node;
            Node* parent = _node->_parent;
            while (parent && current == parent->_right) {
                current = parent;
                parent = parent->_parent;
            }
            _node = parent;
        }
        return *this;
    }

    Self operator--() {
        if (_node == nullptr) {
            Node* last = _root;
            while (last && last->_right) last = last->_right;
            _node = last;
        } else if (_node->_left) {
            Node* max = _node->_left;
            while (max->_right) max = max->_right;
            _node = max;
        } else {
            Node* current = _node;
            Node* parent = _node->_parent;
            while (parent && current == parent->_left) {
                current = parent;
                parent = parent->_parent;
            }
            _node = parent;
        }
        return *this;
    }
};

template<class K, class T, class KeyOfT>
class RBTree {
private:
    RBTreeNode<T>* _root = nullptr;
    typedef RBTreeNode<T> Node;

    int _Height(Node* root) {
        if (root == nullptr) return 0;
        int leftHeight = _Height(root->_left);
        int rightHeight = _Height(root->_right);
        return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
    }

    int _Size(Node* root) {
        if (root == nullptr) return 0;
        return _Size(root->_left) + _Size(root->_right) + 1;
    }

public:
    typedef RBTreeIterator<T, T&, T*> Iterator;
    typedef RBTreeIterator<T, const T&, const T*> ConstIterator;

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

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

    Node* Find(const K& key) {
        Node* curr = _root;
        while (curr) {
            if (KeyOfT()(curr->_data) < key) curr = curr->_right;
            else if (KeyOfT()(curr->_data) > key) curr = curr->_left;
            else return curr;
        }
        return nullptr;
    }

    pair<Iterator, bool> Insert(const T& data) {
        if (_root == nullptr) {
            _root = new Node(data);
            _root->_col = BLACK;
            return { Iterator(_root, _root), true };
        }

        Node* current = _root;
        Node* parent = nullptr;
        KeyOfT kot;

        while (current) {
            if (kot(current->_data) < kot(data)) {
                parent = current;
                current = current->_right;
            } else if (kot(current->_data) > kot(data)) {
                parent = current;
                current = current->_left;
            } else {
                return { Iterator(current, _root), false };
            }
        }

        Node* newNode = new Node(data);
        newNode->_col = RED;
        if (kot(data) < kot(parent->_data)) parent->_left = newNode;
        else parent->_right = newNode;

        newNode->_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 = BLACK;
                    uncle->_col = BLACK;
                    grandfather->_col = RED;
                    current = grandfather;
                    parent = grandfather->_parent;
                } else {
                    if (current == parent->_left) {
                        RotateR(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    } else {
                        RotateL(parent);
                        RotateR(grandfather);
                        current->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    break;
                }
            } else {
                Node* uncle = grandfather->_left;
                if (uncle && uncle->_col == RED) {
                    parent->_col = BLACK;
                    uncle->_col = BLACK;
                    grandfather->_col = RED;
                    current = grandfather;
                    parent = grandfather->_parent;
                } else {
                    if (current == parent->_right) {
                        RotateL(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    } else {
                        RotateR(parent);
                        RotateL(grandfather);
                        current->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    break;
                }
            }
        }
        _root->_col = BLACK;
        return { Iterator(newNode, _root), true };
    }

    void RotateR(Node* parent) {
        Node* subL = parent->_left;
        Node* subLR = parent->_left->_right;
        Node* pParent = parent->_parent;

        parent->_left = subLR;
        if (subLR) subLR->_parent = parent;

        parent->_parent = subL;
        subL->_right = parent;

        if (parent == _root) {
            _root = subL;
            subL->_parent = nullptr;
        } else {
            if (parent == pParent->_left) pParent->_left = subL;
            else pParent->_right = subL;
            subL->_parent = pParent;
        }
    }

    void RotateL(Node* parent) {
        Node* subR = parent->_right;
        Node* subRL = parent->_right->_left;
        Node* pParent = parent->_parent;

        parent->_right = subRL;
        if (subRL) subRL->_parent = parent;

        parent->_parent = subR;
        subR->_left = parent;

        if (parent == _root) {
            _root = subR;
            subR->_parent = nullptr;
        } else {
            if (parent == pParent->_left) pParent->_left = subR;
            else pParent->_right = subR;
            subR->_parent = pParent;
        }
    }

    int Height() { return _Height(_root); }
    int Size() { return _Size(_root); }
};

Myset.h

#pragma once
#include "RBTree.h"
namespace mySpace {
    template<class K>
    class set {
    private:
        struct SetKeyOfT {
            const K& operator()(const K& key) { return key; }
        };
        RBTree<K, const K, SetKeyOfT> _t;
    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); }
    };
}

Mymap.h

#pragma once
#include "RBTree.h"
namespace mySpace {
    template<class K, class V>
    class map {
    private:
        struct MapKeyOfT {
            const K& operator()(const pair<K, V>& kv) { return kv.first; }
        };
        RBTree<K, pair<const K, V>, MapKeyOfT> _t;
    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); }
        V& operator[](const K& key) {
            pair<iterator, bool> ret = insert({ key, V() });
            return ret.first->second;
        }
    };
}

测试文件：Test.cpp

#define _CRT_SECURE_NO_WARNINGS 1
#include "Myset.h"
#include "Mymap.h"

void Print(const mySpace::set<int>& s) {
    mySpace::set<int>::const_iterator it = s.end();
    while (it != s.begin()) {
        --it;
        cout << *it << " ";
    }
    cout << endl;
}

int main() {
    cout << "--------------------测试 set 的功能--------------------" << endl;
    mySpace::set<int> s;
    s.insert(5);
    s.insert(1);
    s.insert(3);
    s.insert(2);
    s.insert(6);

    cout << "使用迭代器进行遍历 set 容器中的元素" << endl;
    mySpace::set<int>::iterator sit = s.begin();
    while (sit != s.end()) {
        cout << *sit << " ";
        ++sit;
    }
    cout << endl;

    cout << "反向打印 set 容器中的元素" << endl;
    Print(s);

    cout << "--------------------测试 map 的功能--------------------" << endl;
    mySpace::map<string, string> dict;
    dict.insert({ "sort", "排序" });
    dict.insert({ "left", "左边" });
    dict.insert({ "right", "右边" });
    for (auto& kv : dict) {
        cout << kv.first << ":" << kv.second << endl;
    }

    cout << "修改 left + 插入 insert、string + 将所有的键值对的的值都加上一个！之后 map 容器中的内容是：" << endl;
    dict["left"] = "左边，修改";
    dict["insert"] = "插入";
    dict["string"];

    mySpace::map<string, string>::iterator mit = dict.begin();
    while (mit != dict.end()) {
        mit->second += '!';
        cout << mit->first << ":" << mit->second << endl;
        ++mit;
    }
    cout << endl;
    return 0;
}

运行结果

基本操作

一、前置++操作

1. 本质

红黑树的前置 ++ 操作对应中序遍历的下一个节点（后继节点），实现逻辑基于中序遍历的顺序规则。中序遍历顺序为：左子树 → 根节点 → 右子树。 对于红黑树中的任意节点，其后继节点遵循以下规则：

• 若当前节点有右子树：后继节点是右子树的最左节点（即：右子树中的最小值）。
• 若当前节点无右子树：后继节点是第一个向上找到的父节点，且当前节点位于该父节点的左子树。若找不到这样的父节点，则后继为 end()。

2. 步骤

• 情况 1：当前结点的右子树'不为空' 说明当前结点已访问完毕，下一个要访问的是右子树的中序第一个结点（即右子树的最左结点），直接取右子树的最左结点即可。
• 情况 2：当前结点的右子树'为空' 说明当前结点及所在子树已访问完毕，下一个结点需要沿着'当前结点到根'的祖先路径向上找。
  • 若当前结点是'父结点的左孩子'：根据中序逻辑，下一个访问的结点就是父结点。
  • 若当前结点是'父结点的右孩子'：说明当前子树已访问完，需继续向上找祖先，直到找到一个'当前节点作为其左子节点的祖先节点'。

3. 图示

4. 解释

• 疑问 1：为什么右子树存在时，后继是右子树的最左节点？ 根据中序遍历规则，访问完当前节点后，应访问其右子树，而右子树中最先被访问的节点是其最左节点（最小值）。
• 疑问 2：为什么右子树不存在时，要找'作为左子节点的祖先'？ 当中序遍历访问完某个节点且其右子树为空时，意味着该节点及其子树已遍历完毕，需要回溯到最近的一个尚未完全遍历的父节点。若当前节点是其父节点的右子节点，则说明父节点及其右子树已遍历完，需继续向上回溯直到找到一个'作为左子节点'的祖先。
• 疑问 3：如何处理 end() 情况？ 当向上回溯到根节点仍未找到符合条件的祖先时，parent 会变为 nullptr。此时将 _node 设为 nullptr，表示迭代器已到达 end() 位置。

二、前置–操作

1. 本质

红黑树的前置 -- 操作对应中序遍历的前一个节点（前驱节点），其实现逻辑与前置 ++ 完全对称，但方向相反。 对于红黑树中的任意节点，其前驱节点遵循以下规则：

• 若当前节点有左子树：前驱节点是左子树的最右节点（即：左子树中的最大值）。
• 若当前节点无左子树：前驱节点是第一个向上找到的父节点，且当前节点位于该父节点的右子树。若找不到这样的父节点，则前驱为 rend()。

2. 步骤

• 情况 1：当前结点的左子树'不为空' 这表明当前结点之前，还有可访问的结点，下一个要访问的是左子树的中序最后一个结点（即左子树的最右结点），直接定位到左子树的最右结点即可。
• 情况 2：当前结点的左子树'为空' 意味着当前结点及所在子树已回溯完毕，下一个结点需要沿着'当前结点到根'的祖先路径向上找。
  • 若当前结点是'父结点的右孩子'：按照中序逻辑，下一个要访问的结点就是父结点。
  • 若当前结点是'父结点的左孩子'：说明当前子树已回溯完，得继续向上找祖先，直至找到一个'当前节点作为其右子节点的祖先节点'。

3. 图示

4. 解释

• 疑问 1：为什么左子树存在时，前驱是左子树的最右节点？ 根据中序遍历规则，访问当前节点前，应先访问其左子树，而左子树中最后被访问的节点是其最右节点（最大值）。
• 疑问 2：为什么左子树不存在时，要找'作为右子节点的祖先'？ 当中序遍历访问完某个节点且其左子树为空时，意味着需要回溯到最近的一个尚未完全遍历的父节点。若当前节点是其父节点的左子节点，则说明父节点及其左子树已遍历完，需继续向上回溯直到找到一个'作为右子节点'的祖先。
• 疑问 3：如何处理 rend() 情况？ 前置 -- 应返回树的最右节点（即中序遍历的最后一个节点）。此时需特殊处理：从根节点开始，一路向右找到最右节点。若树为空，--end() 仍返回 end()。

情况 1：当迭代器指向 end()（即 _node 为 nullptr）时

if (_node == nullptr) {
    Node* last = _root;
    while (last && last->_right) last = last->_right;
    _node = last;
}

代码解释

片段一：仿函数的设计

由于红黑树（RBTree）采用泛型设计，无法直接判断模板参数 T 具体是单纯的键类型 K（如：set 的场景），还是键值对类型 pair<K, V>（如：map 的场景）。这会导致一个问题：在 insert 逻辑里进行'节点比较'时，默认的比较规则无法满足需求，因为 pair 的默认比较会同时涉及 key 和 value，但我们实际需要只比较 key。 为解决这个问题，我们在 map 和 set 这两个容器层，分别实现了仿函数 MapKeyOfT 和 SetKeyOfT，并将它们传递给红黑树的 KeyOfT 模板参数。这样，红黑树内部就能通过 KeyOfT 仿函数先从 T 类型对象中提取出 key，再用这个 key 进行比较，从而实现'仅按 key 排序 / 插入'的逻辑。

// RBTree.h
template<class K, class T, class KeyOfT>
class RBTree {
    // ...
};

// Myset.h
template<class K>
class set {
    struct SetKeyOfT {
        const K& operator()(const K& key) { return key; }
    };
    RBTree<K, const K, SetKeyOfT> _t;
};

// Mymap.h
template<class K, class V>
class map {
    struct MapKeyOfT {
        const K& operator()(const pair<K, V>& kv) { return kv.first; }
    };
    RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};

片段二：迭代器的设计

这里的 iterator 的实现思路与 list 的 iterator 框架一致：用一个类型封装'结点指针'，再通过重载运算符，让迭代器能像指针一样完成访问行为（如：*it、++it 等）。

1. begin() 迭代器接口函数

在 map 和 set 的使用场景中，迭代器走的是中序遍历逻辑（左子树 → 根结点 → 右子树）。因此：begin() 返回中序遍历的第一个结点对应的 iterator（即：最左结点的迭代器）。

2. end() 迭代器接口函数

end() 表示'遍历结束的位置'，返回 nullptr 指针。 在红黑树迭代器中，end() 的表示需要结合中序遍历的边界条件处理。STL 源码中，红黑树额外设计了一个哨兵头结点作为 end()，与我们用 nullptr 表示 end() 的思路功能等价，但实现细节不同。

3. --end() 的特殊处理 当对 end() 执行 --（即：--end()）时，若 end() 由 nullptr 表示，需特殊处理：让迭代器直接指向树的最右结点（即：中序遍历的最后一个有效节点）。这样可以保证反向遍历时，从 end() 回退能正确访问到最后一个元素。

3. 迭代器不可修改性的适配

• set 的 iterator：不支持修改 key。因此将红黑树的第二个模板参数设为 const K。即：RBTree<K, const K, SetKeyOfT> _t;
• map 的 iterator：不支持修改 key，但可修改 value。因此将红黑树的第二个模板参数设为 pair<const K, V>。即：RBTree<K, pair<const K, V>, MapKeyOfT> _t;（pair 的第一个参数 const K 保证 key 不可改）。