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

3. 迭代器 struct 定义，封装节点指针，重载运算符模拟指针行为：

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

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

4. 在红黑树中重命名迭代器类型并编写 begin/end：

typedef __RBTree_Iterator<T> iterator;
iterator begin() {
    Node* cur = _root;
    while(cur->_left) cur = cur->_left;
    return iterator(cur);
}
iterator end() { return iterator(nullptr); }

三、红黑树的插入

由于节点存储内容变为 T 类型，插入时无法直接访问 kv.first 进行比较，需引入 KeyOfT 仿函数提取 Key 值。

1. KeyOfT 是一个类模板仿函数，重载 operator() 用于取出节点存储内容中的 Key。
2. 插入逻辑保持不变，但在比较时使用 kot(data) 和 kot(cur->_data)。

template<typename K,typename T,typename KeyOfT>
class RBTree {
public:
    bool Insert(const T& data) {
        if(_root == nullptr) {
            _root = new Node(data);
            _root->_col = BLACK;
            return true;
        } else {
            Node* parent = nullptr;
            Node* cur = _root;
            while(cur) {
                if(kot(data) > kot(cur->_data)) {
                    parent = cur;
                    cur = cur->_right;
                } else if(kot(data) < kot(cur->_data)) {
                    parent = cur;
                    cur = cur->_left;
                } else {
                    return false;
                }
            }
            cur = new Node(data);
            cur->_col = RED;
            cur->_parent = parent;
            if(kot(cur->_data) > kot(parent->_data)) {
                parent->_right = cur;
            } else {
                parent->_left = cur;
            }
            // 旋转与变色调整逻辑...
            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;
                        _root->_col = BLACK;
                    } 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;
                        _root->_col = BLACK;
                    } 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;
                    }
                }
            }
            return true;
        }
    }
private:
    Node* _root;
    KeyOfT kot;
    // 旋转函数省略...
};

四、map/set 的封装

set 的封装

1. 为避免命名冲突，放入自定义命名空间。Set 只需一个 K 作为模板参数。
2. 实现 SetKeyOfT 仿函数，直接返回 data。
3. 迭代器需注意 typename 关键字，告知编译器取出的 iterator 是类型而非变量。

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

map 的封装

1. Map 需要 K 和 V 两个模板参数。
2. 实现 MapKeyOfT 仿函数，返回 pair 的 first 成员。
3. 注意 map 存储类型为 pair<const K, V>，确保 key 不可变。

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

五、测试

测试一

#include<iostream>
#include"MyMap.h"
#include"MySet.h"
using namespace std;
int main() {
    wzx::set<int> s;
    s.Insert(5); s.Insert(2); s.Insert(9); s.Insert(1);
    wzx::set<int>::iterator it = s.begin();
    while(it != s.end()) {
        cout << *it << ' ';
        ++it;
    }
    cout << endl << endl;

    wzx::map<int,int> m;
    m.Insert(make_pair(3,3)); m.Insert(make_pair(1,1));
    m.Insert(make_pair(8,8)); m.Insert(make_pair(6,6));
    wzx::map<int,int>::iterator mit = m.begin();
    while(mit != m.end()) {
        cout << mit->first << ':' << mit->second << endl;
        ++mit;
    }
    cout << endl;
    return 0;
}

测试二

测试发现当前实现允许修改 map 的 key 和 set 的 key，不符合 STL 规范。STL 要求 map 的 key 不可修改，value 可修改；set 的 key 不可修改。

六、学习 STL 库中的实现

set 如何保证 key 不被修改

STL 中 set 的普通迭代器是由 const_iterator 转换而来的。因此，即使使用普通迭代器遍历，也无法修改 key。

map 如何保证 key 不被修改

Map 通过给 pair 中的 key 类型加 const 修饰来实现（pair<const K, V>）。这样既限制了 key 修改，又保留了 value 的可变性。不能像 set 那样完全依赖 const_iterator，否则 value 也会被限制。

八、set 的改进

红黑树的 const 迭代器

通过控制迭代器的模板参数 Ptr 和 Ref，区分普通迭代器和 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){}
    __RBTree_iterator(const Iterator& it):_node(it._node){}
    Ref operator*(){ return _node->_data; }
    Ptr operator->(){ return &_node->_data; }
    // ... 其他运算符
};

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

改进 set

将 set 的 iterator 定义为 const_iterator，并添加 const 版本的 begin/end。

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

九、map 的改进

1. Map 内部存储类型改为 pair<const K, V>，从源头禁止 key 修改。
2. 添加 const 迭代器支持。
3. 实现 operator[] 接口，调用 insert 并返回 value 引用。

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

测试

验证普通迭代器 key 不可改、value 可改；const 迭代器两者均不可改。

十、完善 set 的 Insert

1. 修改返回值以匹配 STL，返回 pair<iterator, bool>。
2. 解决类型不匹配问题：红黑树返回普通迭代器，set 使用 const 迭代器。需在迭代器构造函数中添加接受普通迭代器的重载，进行浅拷贝。

// RBTree 迭代器构造重载
__RBTree_iterator(const Iterator& it):_node(it._node){}

// Set Insert 实现
pair<iterator,bool>Insert(const K& key) {
    pair<typename RBTree<K, K, SetKeyOfT>::iterator,bool> kv = _t.Insert(key);
    return pair<iterator,bool>(kv.first, kv.second);
}

十一、源代码

MyMap.h

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

Myset.h

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

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

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

template<typename K,typename T,typename KeyOfT>
class RBTree {
    typedef RBTreeNode<T> Node;
public:
    RBTree():_root(nullptr){}
    typedef __RBTree_iterator<T, T*, T&> iterator;
    typedef __RBTree_iterator<T,const T*,const T&> const_iterator;
    iterator begin() {
        Node* cur = _root;
        while(cur->_left) cur = cur->_left;
        return iterator(cur);
    }
    iterator end() { return iterator(nullptr); }
    const_iterator begin() const {
        Node* cur = _root;
        while(cur->_left) cur = cur->_left;
        return const_iterator(cur);
    }
    const_iterator end() const { return const_iterator(nullptr); }
    KeyOfT kot;
    Node* Find(const K& key) {
        Node* cur = _root;
        while(cur) {
            if(key > kot(cur->_data)) cur = cur->_right;
            else if(key < kot(cur->_data)) cur = cur->_left;
            else return cur;
        }
        return nullptr;
    }
    pair<iterator,bool>Insert(const T& data) {
        if(_root == nullptr) {
            _root = new Node(data);
            _root->_col = BLACK;
            return make_pair(iterator(_root),true);
        } else {
            Node* parent = nullptr;
            Node* cur = _root;
            while(cur) {
                if(kot(data) > kot(cur->_data)) { parent = cur; cur = cur->_right; }
                else if(kot(data) < kot(cur->_data)) { parent = cur; cur = cur->_left; }
                else { return make_pair(iterator(cur),false); }
            }
            cur = new Node(data);
            cur->_col = RED;
            Node* newnode = cur;
            cur->_parent = parent;
            if(kot(cur->_data) > kot(parent->_data)) parent->_right = cur;
            else parent->_left = cur;
            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;
                        _root->_col = BLACK;
                    } 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;
                        _root->_col = BLACK;
                    } 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;
                    }
                }
            }
            return make_pair(iterator(newnode),true);
        }
    }
    bool CheckColour(Node* root,int blacksum,int leftpathblack) {
        if(root == nullptr) {
            if(blacksum != leftpathblack) { cout << "每条路径上的黑色节点的数目不同" << endl; return false; }
            return true;
        }
        Node* parent = root->_parent;
        if(parent && parent->_col == RED && root->_col == RED) {
            cout << root->_data << ':' << "出现两个连续的红节点" << endl; return false;
        }
        if(root->_col == BLACK) blacksum++;
        return CheckColour(root->_left, blacksum, leftpathblack) && CheckColour(root->_right, blacksum, leftpathblack);
    }
    bool IsBalance() { return IsBalance(_root); }
    bool IsBalance(Node* root) {
        if(root == nullptr) return true;
        if(root->_col != BLACK) { cout << "根节点颜色错误不为黑色" << endl; return false; }
        int leftpathbalck = 0;
        Node* cur = root;
        while(cur) { if(cur->_col == BLACK) leftpathbalck++; cur = cur->_left; }
        return CheckColour(root,0, leftpathbalck);
    }
    int Height() { return Height(_root); }
    int Height(Node* root) {
        if(root == nullptr) return 0;
        int heightL = Height(root->_left);
        int heightR = Height(root->_right);
        return heightL > heightR ? heightL +1 : heightR +1;
    }
private:
    void RotateL(Node* parent) {
        Node* ppnode = parent->_parent;
        Node* cur = parent->_right;
        Node* curleft = cur->_left;
        parent->_right = curleft;
        if(curleft) curleft->_parent = parent;
        cur->_left = parent;
        parent->_parent = cur;
        if(ppnode == nullptr) { _root = cur; cur->_parent = nullptr; }
        else {
            cur->_parent = ppnode;
            if(ppnode->_left == parent) ppnode->_left = cur;
            else ppnode->_right = cur;
        }
    }
    void RotateR(Node* parent) {
        Node* ppnode = parent->_parent;
        Node* cur = parent->_left;
        Node* curright = cur->_right;
        parent->_left = curright;
        if(curright) curright->_parent = parent;
        cur->_right = parent;
        parent->_parent = cur;
        if(ppnode == nullptr) { _root = cur; cur->_parent = nullptr; }
        else {
            if(ppnode->_left == parent) { ppnode->_left = cur; cur->_parent = ppnode; }
            else { ppnode->_right = cur; cur->_parent = ppnode; }
        }
    }
private:
    Node* _root;
};

test.cpp

#include<iostream>
#include"MyMap.h"
#include"MySet.h"
using namespace std;
int main() {
    wzx::map<string, string> m;
    m.Insert(make_pair("insert","xxx"));
    m["string"];
    for(const auto& e : m) {
        cout << e.first << ':' << e.second << endl;
    }
    cout << endl;
    m["string"]="字符串";
    m["insert"]="插入";
    m["hello"]="你好";
    for(const auto& e : m) {
        cout << e.first << ':' << e.second << endl;
    }
    cout << endl;
    return 0;
}

总结

通过上述步骤，我们完成了基于红黑树的 map 和 set 模拟实现。核心在于理解 rb_tree 的泛型设计、迭代器的 const 正确性处理以及插入算法的细节。