C++ 基于红黑树实现 set 和 map 容器

注意源码中第一个模板参数 Key 是传给 find/erase 等函数做形参类型的。对于 set 而言两个参数是一样的，但对于 map 而言 insert 的是 pair 对象，但是 find 和 erase 的是 Key 对象。

2. 修改红黑树

2.1 参数

我们需要对红黑树进行更改以支持不同的搜索场景。相比源码调整一下，key 参数用 K，value 参数用 V，红黑树中的数据类型使用 T。

因为 RBTree 实现了泛型不知道 T 参数到底是 K 还是 pair<K, V>，insert 内部进行比较时无法直接比较，需要任何时候只比较 key。我们在 map 和 set 层分别实现一个 MapKeyOfT 和 SetKeyOfT 的仿函数传给 RBTree 的 KeyOfT，然后 RBTree 中通过 KeyOfT 仿函数取出 T 类型对象中的 key，再进行比较。

对 map 而言，MapKeyOfT 直接返回 pair<key, value> 中的 key：

struct MapKeyOfT {
    const K& operator()(const pair<K, V>& kv) {
        return kv.first;
    }
};

对 set 而言，SetKeyOfT 直接返回 key：

struct SetKeyOfT {
    const K& operator()(const K& key) {
        return key;
    }
};

红黑树的参数类型变为：

template<class K, class T, class KeyOfT>
class RBTree { ... };

2.2 迭代器

set 和 map 的迭代器都是双向迭代器，封装底层红黑树的迭代器来实现。最重要的是 operator++ 和 operator-- 的实现。map 和 set 的迭代器走的是中序遍历（左子树->根结点->右子树）。

• 迭代器++时，如果 it 指向的结点的右子树不为空，找右子树的最左结点。
• 迭代器++时，如果右子树为空，沿着当前结点到根的祖先路径向上找，直到找到孩子是父亲左的那个祖先就是中序要问题的下一个结点。

end() 表示为 nullptr。STL 源码中增加了哨兵位头结点作为 end()，我们这里用 nullptr 充当 end()。

• 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

2.3 map 支持 []

map 要支持 [] 主要需要修改 insert 返回值支持，修改 RBtree 中的 insert 返回值为 pair<Iterator, bool> Insert(const T& data)。

2.4 代码实现

下面是修改后的红黑树的代码实现：

#pragma once
#include<iostream>
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) {}
};

// 迭代器
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; }
    Self& operator++() {
        if (_node->_right) {
            Node* minleft = _node->_right;
            while (minleft->_left) { minleft = minleft->_left; }
            _node = minleft;
        } else {
            Node* cur = _node;
            Node* parent = cur->_parent;
            while (parent && cur == parent->_right) { cur = parent; parent = parent->_parent; }
            _node = parent;
        }
        return *this;
    }
    Self& operator--() {
        if (_node == nullptr) {
            Node* minright = _root;
            while (minright && minright->_right) { minright = minright->_right; }
            _node = minright;
        } else if (_node->_left) {
            Node* minright = _node->_left;
            while (minright->_right) { minright = minright->_right; }
            _node = minright;
        } else {
            Node* cur = _node;
            Node* parent = cur->_parent;
            while (parent && cur == parent->_left) { cur = parent; parent = parent->_parent; }
            _node = parent;
        }
        return *this;
    }
    bool operator!=(const Self& s) { return _node != s._node; }
    bool operator==(const Self& s) { 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*> Const_Iterator;
    Iterator begin() {
        Node* cur = _root;
        while (cur && cur->_left) { cur = cur->_left; }
        return Iterator(cur, _root);
    }
    Iterator end() { return Iterator(nullptr, _root); }
    Const_Iterator begin() const {
        Node* cur = _root;
        while (cur && cur->_left) { cur = cur->_left; }
        return Iterator(cur, _root);
    }
    Const_Iterator end() const { return Iterator(nullptr, _root); }
    pair<Iterator, bool> Insert(const T& data) {
        if (_root == nullptr) {
            _root = new Node(data);
            _root->_col = BLACK;
            return { Iterator(_root, _root), true };
        }
        KeyOfT kot;
        Node* cur = _root;
        Node* parent = nullptr;
        while (cur) {
            if (kot(cur->_data) < kot(data)) {
                parent = cur; cur = parent->_right;
            } else if (kot(cur->_data) > kot(data)) {
                parent = cur; cur = parent->_left;
            } else {
                return { Iterator(cur, _root), false };
            }
        }
        cur = new Node(data);
        Node* newnode = cur;
        cur->_col = RED;
        if (kot(parent->_data) > kot(data)) {
            parent->_left = cur;
        } else {
            parent->_right = 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 = BLACK; 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 = BLACK; 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 { Iterator(newnode, _root), true };
    }
    Iterator Find(const K& key) {
        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();
    }
    bool IsBalance() {
        if (_root == nullptr) return true;
        if (_root->_col == RED) return false;
        int refNum = 0;
        Node* cur = _root;
        while (cur) {
            if (cur->_col == BLACK) { ++refNum; }
            cur = cur->_left;
        }
        return Check(_root, 0, refNum);
    }
private:
    bool Check(Node* root, int blackNum, const int refNum) {
        if (root == nullptr) {
            if (refNum != blackNum) { cout << "存在黑色结点的数量不相等的路径" << endl; return false; }
            return true;
        }
        if (root->_col == RED && root->_parent->_col == RED) {
            cout << root->_data << "存在连续的红色结点" << endl; return false;
        }
        if (root->_col == BLACK) { blackNum++; }
        return Check(root->_left, blackNum, refNum) && Check(root->_right, blackNum, refNum);
    }
    void RotateR(Node* parent) {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        Node* Pparent = parent->_parent;
        parent->_left = subLR;
        if (subLR) subLR->_parent = parent;
        subL->_right = parent;
        parent->_parent = subL;
        if (parent == _root) { _root = subL; subL->_parent = nullptr; }
        else {
            if (parent == Pparent->_left) { Pparent->_left = subL; }
            else { Pparent->_right = subL; }
            subL->_parent = Pparent;
        }
    }
    void RotateL(Node* parent) {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        Node* Pparent = parent->_parent;
        parent->_right = subRL;
        if (subRL) subRL->_parent = parent;
        subR->_left = parent;
        parent->_parent = subR;
        if (parent == _root) { _root = subR; subR->_parent = nullptr; }
        else {
            if (parent == Pparent->_left) { Pparent->_left = subR; }
            else { Pparent->_right = subR; }
            subR->_parent = Pparent;
        }
    }
    Node* _root = nullptr;
};

3. 实现 map 和 set

其实当我们修改完红黑树后，实现 map 和 set 就非常简单了，只需要复用红黑树即可。

3.1 set

#pragma once
#include"RBTree.h"
namespace hcy {
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>::Const_Iterator const_iterator;
    iterator begin() { return _rbtree.begin(); }
    iterator end() { return _rbtree.end(); }
    const_iterator begin() const { return _rbtree.begin(); }
    const_iterator end() const { return _rbtree.end(); }
    pair<iterator, bool> insert(const K& key) { return _rbtree.Insert(key); }
    iterator find(const K& key) { return _rbtree.Find(key); }
private:
    RBTree<K, const K, SetKeyOfT> _rbtree;
};
}

3.2 map

#pragma once
#include"RBTree.h"
namespace hcy {
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>::Const_Iterator const_iterator;
    iterator begin() { return _rbtree.begin(); }
    iterator end() { return _rbtree.end(); }
    const_iterator begin() const { return _rbtree.begin(); }
    const_iterator end() const { return _rbtree.end(); }
    pair<iterator, bool> insert(const pair<K, V>& kv) { return _rbtree.Insert(kv); }
    V& operator[](const K& key) {
        pair<iterator, bool> ret = _rbtree.Insert({ key, V() });
        return ret.first->second;
    }
    iterator find(const K& key) { return _rbtree.Find(key); }
private:
    RBTree<K, pair<const K, V>, MapKeyOfT> _rbtree;
};
}

4. 小结

我们通过自己实现红黑树并且复用它来构建简单的 map 和 set 可以帮助我们：

4.1 深化对数据结构的理解

红黑树是一种复杂的自平衡二叉搜索树，其核心在于通过颜色规则（红 / 黑）和旋转操作维持平衡性，保证插入、删除、查找的时间复杂度稳定在 O(log n)。我们在实现红黑树的过程中，能深入理解平衡树的设计思想、旋转操作的具体逻辑和触发场景、颜色调整的规则等。

而复用红黑树实现 map 和 set 时，能进一步理解这两种容器的底层原理：set 本质是'键即值'的特殊结构，依赖红黑树的键唯一性；map 是'键值对'结构，需在红黑树节点中存储键和值，并基于键进行排序和查找。

4.2 强化'抽象与复用'的编程思维

红黑树是 map 和 set 的'底层引擎'，两者都依赖红黑树的排序和查找能力，仅在节点存储的数据和接口设计上有差异。通过复用红黑树实现这两种容器，能直观体会'抽象复用'的思想：将红黑树的核心逻辑抽象为独立组件，基于该组件通过'适配层'实现不同容器，减少代码冗余。