从红黑树到 map/set：一次完整的 STL 容器模拟实现

从红黑树到 map/set：一次完整的 STL 容器模拟实现

在继续之前，最好你已经知道红黑树的基本操作。这篇笔记记录的是：给定一个能跑的通的红黑树结构，如何把它封装成 std::set 和 std::map 那样的容器。

STL 里的那点模板把戏

看 set 的模板参数，红黑树部分传了两个 Key：

template<class Key, class Compare = less<Key>, class Alloc = alloc>
class set {
    typedef rb_tree<Key, Key, identity<Key>, key_compare, Alloc> rep_type;
    rep_type t;
};

多传一个重复的类型，为的是和 map 的模板参数保持一致，让底层红黑树能复用同一份代码。map 传的是 pair<const Key, T>：

template<class Key, class T, class Compare = less<Key>, class Alloc = alloc>
class map {
    typedef rb_tree<Key, pair<const Key, T>, select1st<value_type>, key_compare, Alloc> rep_type;
    rep_type t;
};

所以红黑树节点实际存的数据是泛化的 T，set 是 Key，map 是 pair。而第一个模板参数 Key 独立出来，方便 find 这类接口定义参数类型。

红黑树迭代器：按中序遍历走

为了支持容器遍历，我们需要给红黑树做一个迭代器。它本质上就是对节点指针的包装。

节点定义我改成这样，_data 存的是 T：

template<typename T>
struct RBTreeNode {
    T _data;
    RBTreeNode<T>* _left;
    RBTreeNode<T>* _right;
    RBTreeNode<T>* _parent;
    Colour _col;
};

然后树的类模板改成 template<typename K, typename T> class RBTree，节点用 RBTreeNode<T>。

迭代器需要实现 * 返回数据引用，-> 返回数据地址，以及递增和递减按照中序遍历逻辑移动。

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

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

红黑树内部提供 begin() 和 end()。begin 是最左节点，end 直接是 nullptr，用一个空迭代器表示结尾。这里 end() 的迭代器 _node 是空，一般没什么问题，不过要注意不能对 end 解引用。

插入时怎么从 T 里取出 Key

节点里只存了 T，可是红黑树插入时要比较大小，就必须能从 T 提取出 Key。这里的解决方案是加一个模板仿函数 KeyOfT。对于 set，T 就是 Key，仿函数直接返回本身；对于 map，T 是 pair，仿函数返回 first。

树的模板变成 template<typename K, typename T, typename KeyOfT> class RBTree。插入逻辑的关键部分：

bool Insert(const T& data) {
    if (_root == nullptr) {
        _root = new Node(data);
        _root->_col = BLACK;
        return true;
    }
    // ... 按照 kot(data) 比较大小找插入位置
    while (parent && parent->_col == RED) {
        // 变色与旋转逻辑...
    }
    return true;
}
KeyOfT kot;

这样树里面调用 kot(data) 拿到 key，完全不用管 T 是啥。

封装 set 和 map

有了上面这些东西，封装就很简单了。我用一个命名空间 wzx 放自己的实现。

set 模板：

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 模板，注意 pair 里的 key 要加 const，避免外部不小心修改它：

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

set 的 key 不允许修改，map 的 key 也不能修改但 value 可以。为了实现这个约束，set 的迭代器应该要返回 const 引用。目前上面的代码返回的是普通引用，这会导致可以修改 set 的元素。标准库的做法是让 set 的 iterator 和 const_iterator 都是常量迭代器。后面会说怎么处理。

让 set 的 key 真正不可变：const 迭代器

为了让 set 的项不可修改，最简单的方式是把 set 里的 iterator 和 const_iterator 都弄成常量迭代器。也就是说，从红黑树那里拿到的迭代器，不管普通版还是 const 版，解引用都返回 const T&。

迭代器模板可以进一步改造，多引入两个模板参数 Ptr 和 Ref：

template<typename T, typename Ptr, typename Ref>
struct __RBTree_iterator {
    // ...
    Ref operator*() { return _node->_data; }
    Ptr operator->() { return &_node->_data; }
};

然后在红黑树里定义两种迭代器类型：

typedef __RBTree_iterator<T, T*, T&> iterator;
typedef __RBTree_iterator<T, const T*, const T&> const_iterator;

set 里只需要把公开的 iterator 名字直接映射成 const_iterator，禁止通过迭代器修改：

typedef typename RBTree<K, K, SetKeyOfT>::const_iterator iterator;
typedef typename RBTree<K, K, SetKeyOfT>::const_iterator const_iterator;

这样 set 的 begin 和 end 返回的都是 const_iterator，自然就只读。

map 因为 pair 里 key 已经加了 const，所以普通迭代器也能保证 key 不可变，但 value 可变，因此 map 的普通迭代器还是返回非 const 引用。

有一种偷懒的办法是不用改迭代器模板，直接在 set 里把 iterator 变成 const_iterator 的别名，但需要红黑树提供 const_iterator 成员函数。这里为了通用性，我还是改了迭代器模板，这样控制更灵活。

insert 返回值与 operator[]

标准 map::insert 的返回值是 pair<iterator, bool>，可以知道插入是否成功和位置。我们需要修改红黑树 Insert 方法，在插入新节点后返回一个包含迭代器和布尔值的 pair。同时，因为旋转调整可能会改变树结构，但要返回的迭代器始终指向新插入的节点，所以我们在一开始 new 出来节点后就保存起来。

pair<iterator, bool> Insert(const T& data) {
    if (_root == nullptr) {
        _root = new Node(data);
        _root->_col = BLACK;
        return make_pair(iterator(_root), true);
    }
    // ... 找到位置，如果已存在，返回 make_pair(iterator(exist_node), false)
    // 如果不存在， newnode = new Node(data); 然后执行旋转调整
    return make_pair(iterator(newnode), true);
}

有了这个之后，map 的 operator[] 就很短了：

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

注意如果 key 不存在，Insert 会插入一个默认构造的 value，然后返回引用；如果存在，就返回已有的引用。

跑一下测试

写几个简单的测试用例，覆盖插入、遍历，以及修改限制的编译检查。比如：

• 遍历 set，尝试修改 *it，编译器应该报错。
• 遍历 map，修改 it->second 可以，修改 it->first 报错。
• map 的 operator[] 可以用做插入和更新。

代码不贴了，就是常规的遍历打印。至此，一个基本可用的、基于红黑树的 set 和 map 就算模拟完成了。还有很多细节没处理（比如拷贝构造、删除、迭代器失效），但核心的骨架已经能跑起来。