STL 容器底层原理：基于红黑树模拟实现 map 与 set

二、红黑树的迭代器

为了实现 map/set 的遍历，我们需要模拟实现红黑树的迭代器。原红黑树设计为 key_value 结构，现在需改进为通用结构 template<typename K, typename T>，其中 T 表示节点实际存储的类型。

1. 节点定义

节点中存储的内容统一使用 _data 成员变量，类型为 T。

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

2. 迭代器实现

迭代器封装了节点指针，重载运算符以模拟指针行为。注意区分普通迭代器和 const 迭代器的返回类型。

template<typename T, typename Ptr, typename Ref>
struct __RBTree_iterator {
    typedef RBTreeNode<T> Node;
    typedef __RBTree_iterator<T, Ptr, Ref> Self;
    typedef __RBTree_iterator<T, T*, T&> Iterator;
    Node* _node;

    __RBTree_iterator(Node* root) : _node(root) {}
    // 允许从普通迭代器构造 const 迭代器
    __RBTree_iterator(const Iterator& it) : _node(it._node) {}

    Ref operator*() { return _node->_data; }
    Ptr 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;
    }

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

在红黑树类中，分别 typedef 出 iterator 和 const_iterator：

template<typename K, typename T, typename KeyOfT>
class RBTree {
    typedef RBTreeNode<T> Node;
public:
    typedef __RBTree_iterator<T, T*, T&> iterator;
    typedef __RBTree_iterator<T, const T*, const T&> const_iterator;
    // ... begin/end 实现略
};

三、红黑树的插入

由于节点存储的是通用类型 T，插入时无法直接比较 key。我们需要引入第三个模板参数 KeyOfT（仿函数），用于提取节点数据中的 key 值进行比较。

template<typename K, typename T, typename KeyOfT>
class RBTree {
    KeyOfT kot; // 仿函数对象
public:
    bool Insert(const T& data) {
        if (_root == nullptr) {
            _root = new Node(data);
            _root->_col = BLACK;
            return true;
        }
        // ... 查找位置并插入逻辑
        // 比较时使用 kot(data) 和 kot(cur->_data)
        // ... 旋转与变色调整
        return true;
    }
};

四、map/set 的封装

1. set 的封装

set 只需一个模板参数 K。仿函数 SetKeyOfT 直接返回 data。为了防止 key 被修改，迭代器底层使用 const_iterator。

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>::const_iterator iterator;
    typedef typename RBTree<K, K, SetKeyOfT>::const_iterator const_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 是 key_value 结构。仿函数 MapKeyOfT 取出 pair 的 first。为了保证 key 不可修改，pair 中的 key 定义为 const。

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;
    typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::const_iterator 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<const K, V> kv) { return _t.Insert(kv); }
    V& operator[](const K& key) {
        pair<iterator, bool> kv = Insert(make_pair(key, V()));
        return kv.first->second;
    }
private:
    RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};
}

五、测试验证

1. 基本功能测试

int main() {
    wzx::set<int> s;
    s.Insert(5); s.Insert(2); s.Insert(9); s.Insert(1);
    for (auto it = s.begin(); it != s.end(); ++it) {
        cout << *it << ' ';
    }
    cout << endl;

    wzx::map<string, string> m;
    m["insert"] = "xxx";
    m["string"] = "test";
    for (const auto& e : m) {
        cout << e.first << ':' << e.second << endl;
    }
    return 0;
}

2. 限制 key 修改测试

尝试通过迭代器修改 set 的 key 或 map 的 key，编译器将报错，证明 const 修饰生效。map 的 value 仍可正常修改。

六、总结

通过上述步骤，我们完成了基于红黑树的 map 和 set 模拟实现。核心点在于：

1. 泛型设计：利用 KeyOfT 仿函数解耦存储类型与比较键。
2. 迭代器控制：通过 const_iterator 确保 key 的不可变性。
3. 接口完善：实现了 insert 返回值优化及 operator[] 语法糖。

这套实现逻辑与 STL 标准库高度一致，有助于深入理解 C++ 关联式容器的底层机制。