C++ STL 容器：基于红黑树模拟实现 map 与 set

template<typename T>
struct __RBTree_iterator {
    typedef RBTreeNode<T> Node;
    typedef __RBTree_iterator<T> Self;
    Node* _node;

    __RBTree_iterator(Node* node) : _node(node) {}
    
    T& operator*() { return _node->_data; }
    T* operator->() { return &_node->_data; }
    bool operator!=(const Self& it) { return _node != it._node; }

    Self& operator++() {
        Node* cur = _node;
        Node* parent = cur->_parent;
        if (cur->_right) {
            Node* RightMin = cur->_right;
            while (RightMin->_left) RightMin = RightMin->_left;
            _node = RightMin;
        } else {
            while (parent && cur == parent->_right) {
                cur = parent;
                parent = parent->_parent;
            }
            _node = parent;
        }
        return *this;
    }
    // operator-- 逻辑类似，省略部分代码以保持简洁
};

三、红黑树的插入与 KeyOfValue 仿函数

在通用红黑树中，节点存储的是 T 类型，但插入时需要比较 Key 值。由于 T 可能是 Key 也可能是 pair<Key, Value>，我们需要一个机制提取 Key。

STL 引入了第三个模板参数 KeyOfValue（此处称为 KeyOfT），它是一个仿函数对象，负责从存储数据中提取 Key。

template<typename K, typename T, typename KeyOfT>
class RBTree {
    KeyOfT kot; // 仿函数对象
    // 插入逻辑中使用 kot(data) 获取 Key 进行比较
    // ... 旋转与变色逻辑保持不变
};

这样，map 和 set 在实例化时传入不同的仿函数，即可适配各自的存储结构。

四、map 与 set 的封装

1. Set 的封装

set 只需要存储 Key，且 Key 不可修改。我们定义 SetKeyOfT 仿函数，直接返回 Key 的引用。

namespace wzx {
template<typename K>
class set {
private:
    struct SetKeyOfT {
        const K& operator()(const K& data) const { return data; }
    };
public:
    typedef typename RBTree<K, K, SetKeyOfT>::iterator iterator;
    iterator begin() { return _t.begin(); }
    iterator end() { return _t.end(); }
    bool Insert(const K& key) { return _t.Insert(key); }
private:
    RBTree<K, K, SetKeyOfT> _t;
};
}

2. Map 的封装

map 存储 pair<const K, V>。MapKeyOfT 仿函数提取 pair 的 first 成员。

namespace wzx {
template<typename K, typename V>
class map {
private:
    struct MapKeyOfT {
        const K& operator()(const pair<K, V>& kv) const { return kv.first; }
    };
public:
    typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::iterator iterator;
    iterator begin() { return _t.begin(); }
    iterator end() { return _t.end(); }
    bool Insert(const pair<const K, V> kv) { return _t.Insert(kv); }
private:
    RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};
}

五、Const 正确性与 Key 保护

STL 要求 set 的 Key 不可修改，map 的 Key 不可修改但 Value 可修改。

1. Set: 通过将普通迭代器定义为 const_iterator 的重命名，确保解引用后返回 const K&，从而禁止修改。
2. Map: 直接在 pair 的 Key 类型上添加 const 修饰（即 pair<const K, V>）。这样即使通过迭代器访问，Key 也是只读的，而 Value 保持可变。

六、operator[] 与 Insert 返回值优化

1. Insert 返回值

STL 的 insert 返回 pair<iterator, bool>，指示是否插入成功及位置。我们需要改造红黑树的 Insert 方法以匹配此行为。

pair<iterator, bool> Insert(const T& data) {
    // 若根为空，创建新节点，返回 true
    // 若已存在，返回 false
    // 若插入成功，保存新节点指针，调整颜色后返回 true
}

2. operator[] 实现

operator[] 通常用于通过 Key 访问 Value。若 Key 不存在，则插入默认构造的 Value。

V& operator[](const K& key) {
    pair<iterator, bool> kv = Insert(make_pair(key, V()));
    return kv.first->second;
}

七、测试验证

编写测试用例验证以下场景：

1. 基本插入与遍历。
2. 尝试修改 Key 值（应编译失败或逻辑被阻止）。
3. 修改 Value 值（应允许）。
4. operator[] 的自动插入行为。

经过多轮测试，封装后的容器在功能上与 STL 标准库保持一致。

八、核心源代码

以下为简化后的核心头文件结构，包含红黑树、Map 和 Set 的定义。

RBTree.h

#pragma once
#include <iostream>
using namespace std;
enum Colour { RED, BLACK };

template<typename 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) {}
};
// ... 迭代器与 RBTree 类实现 ...

MySet.h

#pragma once
#include "RBTree.h"
namespace wzx {
template<typename K>
class set {
    // ... 见上文封装代码 ...
};
}

MyMap.h

#pragma once
#include "RBTree.h"
namespace wzx {
template<typename K, typename V>
class map {
    // ... 见上文封装代码 ...
};
}

总结

通过本文的逐步推导，我们完成了基于红黑树的 map 和 set 模拟实现。关键点在于利用模板元编程和仿函数解决泛型存储问题，并通过 const 修饰保证数据的完整性。这一过程不仅加深了对 STL 容器底层原理的理解，也锻炼了 C++ 模板编程的能力。