C++ 红黑树封装 set 与 map 底层实现

• 通过框架分析，源码中 rb_tree 用了巧妙的泛型思想实现。
• set 实例化 rb_tree 时第二个模板参数给的是 key，map 实例化时给的是 pair<const key, T>。
• 源码中的 value_type 是红黑树结点中存储的真实数据类型。
• 第一个模板参数 Key 用于 find/erase 等函数做形参类型。

2 模拟实现 map 和 set

2.1 实现出复用红黑树的框架，并支持 insert

参考源码框架，map 和 set 复用之前实现的红黑树。相比源码调整一下，key 参数用 K，value 参数用 V，红黑树中的数据类型使用 T。因为 RBTree 实现了泛型不知道 T 参数导致是 K 还是 pair<K, V>，insert 内部进行比较时无法直接比较，需要我们在 map 和 set 层分别实现一个 MapKeyofpair 和 SetKeyofpair 的仿函数传给 RBtree 的 KeyOfPair。

// map.h
template<class T1, class T2>
bool operator<(const pair<T1, T2>& lhs, const pair<T1, T2>& rhs) {
    return lhs.first < rhs.first || (!(rhs.first < lhs.first) && lhs.second < rhs.second);
}

template<class K, class V>
class map {
    struct MapKeyofpair {
        const K& operator()(const pair<K, V>& kv) { return kv.first; }
    };
public:
    bool insert(const pair<K, V>& kv) {
        return _t.Insert(kv); // 底层是红黑树，调用红黑树的 insert 就行了
    }
private: 
    RBTree<K, pair<K, V>, MapKeyofpair> _t;
};

// set.h
template<class K>
class set {
    struct SetKeyofT {
        const K& operator()(const K& key) { return key; }
    };
public:
    bool insert(const K& key) {
        return _t.Insert(key);
    }
private: 
    RBTree<K, K, SetKeyofpair> _t;
};

// RBtree.h
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) {}
};

template<class K, class T, class KeyOfT>
class RBtree {
private:
    typedef RBTreeNode<T> Node;
    Node* _root = nullptr;
public:
    bool Insert(const T& data) {
        if (_root == nullptr) {
            _root = new Node(data);
            _root->_col = BLACK;
            return 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 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;
        // ... 平衡处理逻辑省略
        return true;
    }
};

2.2 支持 iterator 的实现

红黑树的迭代器不是 ++ 就能得到下一个结点的。Iterator 实现的框架跟 list 的 iterator 思路一致，用一个类型封装结点的指针再通过重载运算符实现。

2.2.1 红黑树迭代器结构

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

2.2.2 迭代器++

迭代器++的核心逻辑只看局部，考虑当前中序局部要访问的下一个结点（左子树 -> 根结点 -> 右子树）。

• 情况一（右子树不为空）：找右子树的最左结点。
• 情况二（右子树为空）：沿着祖先路径向上找，直到找到孩子是父亲左的那个祖先。

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

end() 实现：当 it 指向 50 时，++it 找不到孩子是父亲左的祖先，父指针为空，将 it 中的结点指针置为 nullptr 充当 end。

2.2.4 iterator--

迭代器--的实现跟++的思路完全类似，逻辑正好反过来（右子树 -> 根结点 -> 左子树）。

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

3 注意须知 [实现 map/set]

• set 的 iterator 不支持修改，把 set 的第二个模板参数改成 const K 即可：RBTree<K, const K, SetKeyOfT> _t;
• map 的 iterator 不支持修改 key 但是可以修改 value，把 map 的第二个模板参数 pair 的第一个参数改成 const K 即可：RBTree<K, pair<const K, V>, MapKeyOfT> _t;

3.1 map[] 实现

map 要支持 [] 主要需要修改 insert 返回值支持：

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

3.2 代码实现

// map.h
template<class K, class V>
class map {
    struct mapKeyofpair {
        const K& operator()(const pair<K, V>& kv) { return kv.first; }
    };
public:
    typedef typename RBtree<K, pair<const K, V>, mapKeyofpair>::Iterator iterator;
    typedef typename RBtree<K, pair<const K, V>, mapKeyofpair>::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); }
private: 
    RBtree<K, pair<const K, V>, mapKeyofpair> _t;
};

// set.h
template<class K>
class set {
    struct setKeyofpair {
        const K& operator()(const K& key) { return key; }
    };
public:
    typedef typename RBtree<K, const K, setKeyofpair>::Iterator iterator;
    typedef typename RBtree<K, const K, setKeyofpair>::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& data) { return _t.insert(data); }
private: 
    RBtree<K, const K, setKeyofpair> _t;
};

// RBtree.h 核心逻辑包含旋转、变色、插入等完整实现...