C++封装红黑树实现mymap和myset完整代码

一、源码及框架分析

如果你曾好奇STL中的map和set为何性能如此高效，其奥秘很大程度上藏在SGI-STL 3.0版本的源码里。一个精妙的设计是，它们底层复用了同一棵红黑树（rb_tree）。相关的核心代码，主要分布在stl_tree.h、stl_map.h和stl_set.h这几个文件中。

1.1 框架核心代码

先来看看set和map最简化的骨架定义，其核心思路一目了然：

// stl_set.h
template , class Alloc = alloc>
class set {
public:
    typedef Key key_type;
    typedef Key value_type;
private:
    typedef rb_tree, key_compare, Alloc> rep_type;
    rep_type t; // 底层红黑树
};
// stl_map.h
template , class Alloc = alloc>
class map {
public:
    typedef Key key_type;
    typedef T mapped_type;
    typedef pair value_type;
private:
    typedef rb_tree, key_compare, Alloc> rep_type;
    rep_type t; // 底层红黑树
};

看到了吗？两者内部都定义了一个rep_type，它正是底层红黑树的类型。区别在于，set的value_type就是Key本身，而map的value_type是pair。这个差异，直接决定了后续红黑树需要如何设计。

1.2 红黑树的泛型设计

那么，一棵红黑树如何能同时服务于存储单一键的set和存储键值对的map呢？答案在于巧妙的模板参数设计。首先，红黑树的结点定义如下：

template 
struct __rb_tree_node : public __rb_tree_node_base {
    typedef __rb_tree_node* link_type;
    Value value_field; // 存储的实际数据
};

关键来了，红黑树类本身的模板声明是这样的：

template 
class rb_tree {
    // ...
};

这里有三个模板参数至关重要：

• Key：键的类型。它决定了find、erase等接口接受什么类型的参数。
• Value：结点中实际存储的数据类型。对set，它就是Key；对map，它就是pair。
• KeyOfValue：这是一个仿函数，它的使命是从Value中提取出用于比较的Key。因为红黑树在内部进行查找、插入时，比较的始终是键，而非整个存储值。

1.3 为何需要两个模板参数Key和Value？

你可能会问，既然set的Key和Value相同，为何不合并？而map又为何需要两个？这恰恰是设计的精妙之处。

• Value决定了结点里“装什么”。set装单个键，map装键值对。
• Key决定了外部接口“用什么查”。对于map，虽然内部存的是pair，但用户查找时只需要传入Key类型的关键字即可。

因此，这两个参数分工明确，缺一不可。顺便提一句，源码中的命名（set用Key，map用Key和T，rb_tree用Key和Value）初看可能有些混乱，但其背后的设计思路是非常清晰的。

注：源码中命名风格略有混乱，set 使用 Key，map 使用 Key 和 T，rb_tree 又使用 Key 和 Value，但设计思路清晰。

二、模拟实现 map 和 set

理解了源码的设计哲学，我们自己动手封装一个map和set就水到渠成了。整个过程，其实就是围绕如何设计一个通用的红黑树展开。

2.1 实现红黑树框架，支持插入

第一步，我们先搭建红黑树RBTree的骨架。它的核心是引入一个KeyOfT仿函数，用来从存储的数据T中提取键。

RBTree.h

enum Colour { RED, BLACK };
template
struct RBTreeNode {
    T _data;
    RBTreeNode* _left;
    RBTreeNode* _right;
    RBTreeNode* _parent;
    Colour _col;
    RBTreeNode(const T& data)
        : _data(data), _left(nullptr), _right(nullptr), _parent(nullptr), _col(RED) {}
};
template
class RBTree {
    typedef RBTreeNode Node;
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;
    }
private:
    Node* _root = nullptr;
};

有了这个通用的红黑树，封装map和set就变得异常简单。它们只需要定义各自的“取键”仿函数，然后组合红黑树即可。

Mymap.h

namespace bit {
    template
    class map {
        struct MapKeyOfT {
            const K& operator()(const pair& kv) {
                return kv.first;
            }
        };
    public:
        bool insert(const pair& kv) {
            return _t.Insert(kv);
        }
    private:
        RBTree, MapKeyOfT> _t;
    };
}

Myset.h

namespace bit {
    template
    class set {
        struct SetKeyOfT {
            const K& operator()(const K& key) {
                return key;
            }
        };
    public:
        bool insert(const K& key) {
            return _t.Insert(key);
        }
    private:
        RBTree _t;
    };
}

2.2 支持迭代器

一个容器如果没有迭代器，就失去了灵魂。对于红黑树这样的关联式容器，迭代器的核心在于实现中序遍历的步进逻辑，即operator++和operator--。

迭代器实现思路

• begin()：返回中序遍历的第一个结点，也就是整棵树最左边的结点。
• end()：通常返回nullptr（或像源码那样使用一个哨兵头结点）。
• operator++()：寻找当前结点的“后继”。分两种情况：
  • 如果当前结点有右子树，那么后继就是其右子树中的最左结点。
  • 如果没有右子树，则需要向上回溯，找到第一个“当前结点是其父节点左孩子”的祖先结点。
• operator--()：逻辑与++完全对称，寻找当前结点的“前驱”。

迭代器代码

template
struct RBTreeIterator {
    typedef RBTreeNode Node;
    typedef RBTreeIterator Self;
    Node* _node;
    Node* _root; // 用于处理 end() 的情况
    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* leftMost = _node->_right;
            while (leftMost->_left) leftMost = leftMost->_left;
            _node = leftMost;
        } else {
            Node* cur = _node;
            Node* parent = cur->_parent;
            while (parent && cur == parent->_right) {
                cur = parent;
                parent = cur->_parent;
            }
            _node = parent;
        }
        return *this;
    }
    Self& operator--() {
        if (_node == nullptr) { // 处理 --end()
            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;
    }
    bool operator!=(const Self& s) const { return _node != s._node; }
    bool operator==(const Self& s) const { return _node == s._node; }
};

在 RBTree 中集成迭代器

接下来，需要在红黑树类中提供迭代器的接口：

template
class RBTree {
public:
    typedef RBTreeIterator Iterator;
    typedef RBTreeIterator ConstIterator;
    Iterator Begin() {
        Node* leftMost = _root;
        while (leftMost && leftMost->_left) leftMost = leftMost->_left;
        return Iterator(leftMost, _root);
    }
    Iterator End() { return Iterator(nullptr, _root); }
    // 同理实现 ConstIterator
private:
    Node* _root = nullptr;
};

2.3 支持operator[]

map的[]运算符是其便利性的关键。它的实现巧妙地依赖于insert的返回值。因此，我们首先要修改RBTree::Insert，让它返回一个pair，其中包含插入位置的迭代器和是否插入成功的标志。

pair Insert(const T& data) {
    // 插入逻辑，返回插入位置的迭代器及是否成功
}

这样，在map中实现operator[]就非常简洁了：

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

这行代码堪称经典：尝试插入一个以key为键、V()（值类型的默认构造值）为值的键值对。如果key已存在，insert返回已存在位置的迭代器；如果不存在，则插入新结点并返回其迭代器。最后，无论哪种情况，都返回这个键对应值的引用。完美实现了“查找兼插入”的语义。

2.4 完整代码

将上述所有部分组合起来，就得到了我们完整的map和set模拟实现。以下是最终版的关键头文件概览：

最终版Myset.h

#include "RBTree.h"
namespace bit {
    template
    class set {
        struct SetKeyOfT {
            const K& operator()(const K& key) { return key; }
        };
    public:
        typedef typename RBTree::Iterator iterator;
        typedef typename RBTree::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 insert(const K& key) { return _t.Insert(key); }
        iterator find(const K& key) { return _t.Find(key); }
    private:
        RBTree _t;
    };
}

最终版Mymap.h

#include "RBTree.h"
namespace bit {
    template
    class map {
        struct MapKeyOfT {
            const K& operator()(const pair& kv) { return kv.first; }
        };
    public:
        typedef typename RBTree, MapKeyOfT>::Iterator iterator;
        typedef typename RBTree, MapKeyOfT>::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 insert(const pair& kv) { return _t.Insert(kv); }
        iterator find(const K& key) { return _t.Find(key); }
        V& operator[](const K& key) {
            pair ret = insert(make_pair(key, V()));
            return ret.first->second;
        }
    private:
        RBTree, MapKeyOfT> _t;
    };
}

最终版RBTree.h

完整的RBTree.h需要整合所有上述片段，并补充完整的插入平衡（红黑树的旋转与变色）、查找(Find)、以及迭代器等接口的实现。其核心模块包括：

• 结点(RBTreeNode)定义
• 插入与自平衡逻辑
• 迭代器(RBTreeIterator)的实现
• 对外接口：Begin, End, Find, Insert等

三、总结

通过从头封装红黑树来实现map和set，我们得以深入STL设计的内核。整个过程清晰地展示了泛型编程和代码复用的强大力量：

1. 泛型设计是基石：红黑树通过Value模板参数适配不同的存储内容（单键或键值对），再通过KeyOfT仿函数统一提取比较键，从而实现了高度的通用性。
2. 迭代器是桥梁：中序遍历的步进逻辑是迭代器实现的核心。处理好左右子树为空时向祖先节点的回溯，是正确实现operator++和operator--的关键。
3. operator[]的巧思：map的[]运算符完美利用了insert的返回值，用一行代码实现了查找、插入、访问的三合一功能，既简洁又高效。
4. 权限控制的艺术：通过模板参数巧妙控制修改权限。set通过将Value类型设为const K，使得迭代器解引用得到常量，防止键被修改。map则通过pair直接保护键的常量性，同时允许值被修改。

这种设计不仅极大地减少了代码重复，更体现了面向对象与泛型编程思想结合所带来的优雅与强大。理解它，对于编写高性能、可复用的C++代码至关重要。