C++ 封装红黑树实现 MyMap 与 MySet

C++ 封装红黑树实现 MyMap 与 MySet

在深入关联式容器之前，我们需要理解其底层数据结构。本文基于红黑树（Red-Black Tree），模拟实现 C++ STL 中的 std::map 和 std::set。我们将先梳理红黑树的核心规则与基础操作（插入、删除、旋转），再基于红黑树的底层结构，分别设计 MyMap（适配键值对存储）和 MySet（适配单一键存储）的核心接口。

1. 源码及框架分析

参考 SGI STL 30 版本源代码，map 和 set 的源码主要在 stl_map.h、stl_set.h 和 stl_tree.h 等头文件中。核心结构如下：

// stl_set.h
template<class Key, class Compare = less<Key>, class Alloc = alloc>
class set {
public:
    typedef Key key_type;
    typedef Key value_type;
private:
    // identity<value_type> 用于从 value 中提取 key，因为 set 中 value 就是 key
    typedef rb_tree<key_type, value_type, identity<value_type>, key_compare, Alloc> rep_type; 
    rep_type t; // 红黑树实例
};

// stl_map.h
template<class Key, class T, class Compare = less<Key>, class Alloc = alloc>
class map {
public:
    typedef Key key_type;
    typedef T mapped_type;
    typedef pair<const Key, T> value_type;
private:
    // select1st<value_type> 用于从 pair 中提取第一个元素 (Key)
    typedef rb_tree<key_type, value_type, select1st<value_type>, key_compare, Alloc> rep_type; 
    rep_type t; // 红黑树实例
};

核心设计点

1. 泛型复用：rb_tree 是核心，它不直接写死是存 Key 还是存 Key/Value，而是由第二个模板参数 Value 决定。set 传入 Key，map 传入 pair<const Key, T>。
2. KeyOfValue 仿函数：这是关键。红黑树需要按 Key 排序，但存储的是 Value。因此需要一个仿函数从 Value 中提取 Key。
   • set 使用 identity，直接返回 Value 本身。
   • map 使用 select1st，提取 pair 的 first 成员。
3. 节点结构：

   struct __rb_tree_node_base {
       color_type color;
       base_ptr parent;
       base_ptr left;
       base_ptr right;
   };
   template<class Value>
   struct __rb_tree_node : public __rb_tree_node_base {
       Value value_field; // 存储实际数据
   };
   

2. 模拟实现 Map 和 Set

2.1 封装框架与 KeyOfT

相比源码，我们简化命名，用 K 代表 Key，V 代表 Value，T 代表红黑树内部存储类型。由于 RBTree 不知道 T 是 K 还是 pair<K,V>，我们在 MyMap 和 MySet 层分别实现 MapKeyOfT 和 SetKeyOfT 仿函数传给 RBTree。

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

    // MySet.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, SetKeyOfT> _t;
    };
}

2.2 迭代器实现

迭代器的核心难点在于 operator++ 和 operator--。map 和 set 的迭代器走的是中序遍历（左子树 -> 根结点 -> 右子树）。

• operator++:
  1. 如果当前结点的右子树不为空，下一个结点是右子树的最左结点。
  2. 如果右子树为空，向上找祖先。如果当前结点是父结点的左孩子，则父结点是下一个；如果是右孩子，继续向上找，直到找到某个祖先是其父结点的左孩子为止。
• operator--: 逻辑相反，优先找左子树的最右结点，否则向上找。

注意：end() 通常指向 nullptr 或哨兵节点。这里我们用 nullptr 充当尾后迭代器，并在 --end() 时特殊处理指向最右结点。

2.3 Map 支持 [] 运算符

map 的 [] 运算符本质上是查找，如果不存在则插入默认值并返回引用。这要求 Insert 方法返回 pair<Iterator, bool>。

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

3. 完整代码实现

以下是整合后的红黑树迭代器及核心类实现，包含插入平衡调整与中序遍历逻辑。

#include <iostream>
#include <utility>
using namespace std;

enum Colour { RED, BLACK };

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

    RBTreeNode(const T& data) : _data(data), _col(RED), _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;

    RBTreeIterator(Node* node) : _node(node) {}

    Self& operator++() {
        if (_node->_right) {
            Node* cur = _node->_right;
            while (cur->_left) cur = cur->_left;
            _node = cur;
        } 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) {
            Node* rightMost = _root;
            while (rightMost && rightMost->_right) rightMost = rightMost->_right;
            _node = rightMost;
        } else if (_node->_left) {
            Node* cur = _node->_left;
            while (cur->_right) cur = cur->_right;
            _node = cur;
        } else {
            Node* cur = _node;
            Node* parent = cur->_parent;
            while (parent && cur == parent->_left) {
                cur = parent;
                parent = cur->_parent;
            }
            _node = parent;
        }
        return *this;
    }

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

    Node* _root;
};

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*> ConstIterator;

    Iterator Begin() {
        Node* cur = _root;
        while (cur && cur->_left) cur = cur->_left;
        return Iterator(cur);
    }

    Iterator End() {
        return Iterator(nullptr);
    }

    pair<Iterator, bool> Insert(const T& data) {
        if (!_root) {
            _root = new Node(data);
            _root->_col = BLACK;
            return make_pair(Iterator(_root), 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 make_pair(Iterator(cur), false);
            }
        }

        cur = new Node(data);
        cur->_col = RED;
        if (kot(parent->_data) < kot(data)) {
            parent->_right = cur;
        } else {
            parent->_left = 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 = 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 = 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 make_pair(Iterator(cur), true);
    }

private:
    void RotateL(Node* parent) {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        parent->_right = subRL;
        if (subRL) subRL->_parent = parent;
        Node* parentParent = parent->_parent;
        subR->_left = parent;
        parent->_parent = subR;
        if (!parentParent) _root = subR;
        else {
            if (parent == parentParent->_left) parentParent->_left = subR;
            else parentParent->_right = subR;
        }
        subR->_parent = parentParent;
    }

    void RotateR(Node* parent) {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        parent->_left = subLR;
        if (subLR) subLR->_parent = parent;
        Node* parentParent = parent->_parent;
        subL->_right = parent;
        parent->_parent = subL;
        if (!parentParent) _root = subL;
        else {
            if (parent == parentParent->_left) parentParent->_left = subL;
            else parentParent->_right = subL;
        }
        subL->_parent = parentParent;
    }

    void Destroy(Node* root) {
        if (!root) return;
        Destroy(root->_left);
        Destroy(root->_right);
        delete root;
    }

private:
    Node* _root = nullptr;
};

通过这一过程，不仅能吃透红黑树的平衡机制，更能深刻理解关联式容器'有序、唯一'特性的底层实现逻辑。