手撕 STL 源码：基于红黑树封装 map 与 set

二、模拟实现 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 并返回其引用。

实现步骤：

1. 修改 RBTree 的 Insert，返回 pair<Iterator, bool>。
2. 在 map 的 operator[] 中调用 insert。
3. 根据返回的迭代器访问 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 设计亮点

1. 泛型设计：一棵红黑树通过模板参数适配 set 和 map。
2. 仿函数 KeyOfT：解决了不同类型中取 key 的问题。
3. 迭代器抽象：封装了中序遍历的逻辑。
4. 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 找到最右节点。