封装红黑树实现自定义 set 与 map 容器

End 迭代器的特殊处理

STL 源码中常使用哨兵位（Sentinel Node）来简化边界判断，但这会增加维护成本（如旋转时需更新哨兵指针）。我们可以采用更轻量级的方案：用 nullptr 表示 end() 迭代器。

当执行 end()-- 时，需要特殊处理：此时应返回树中最后一个节点（最右节点）。因此，迭代器成员变量除了当前节点 _node，还需要保存根节点 _root，以便在 nullptr 时能快速定位到最右节点。

普通迭代器与 Const 迭代器

两者的区别仅在于解引用运算符 * 和 -> 的返回值类型。我们可以通过增加两个模板参数 Ref 和 Ptr 来区分它们。普通迭代器对应 T& 和 T*，Const 迭代器对应 const T& 和 const T*。

template<class T, class Ref, class Ptr>
struct RBTreeIterator {
    typedef RBTreeNode<T> Node;
    typedef RBTreeIterator<T, Ref, Ptr> Self;
    
    Node* _root;
    Node* _node;

    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++(int) {
        Self tmp = *this;
        ++(*this);
        return tmp;
    }

    // 前置 --
    Self& operator--() {
        if (_node == nullptr) {
            // end-- 指向树中最后一个节点
            Node* maxRight = _root;
            while (maxRight && maxRight->_right)
                maxRight = maxRight->_right;
            _node = maxRight;
        } 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;
    }

    // 后置 --
    Self& operator--(int) {
        Self tmp = *this;
        --(*this);
        return tmp;
    }

    bool operator==(const Self& s) const { return _node == s._node; }
    bool operator!=(const Self& s) const { return _node != s._node; }
};

Set 与 Map 的关键细节

保证 Key 不可变

set 中的元素即 Key，天然不可变。map 中存储的是 pair<K, V>，为了保证 Key 不被修改，同时允许 Value 被修改，STL 采用了巧妙的技巧：在 map 的红黑树底层存储类型中，将 Key 设为 const，即 pair<const K, V>。

这样，通过 const_iterator 访问时，Key 部分受到保护，而 Value 部分依然可写。在 myset 的实现中，直接将第二个模板参数设为 const K 即可达到同样效果。

operator[] 的实现

map 的 operator[] 功能包含插入、修改和查找。其底层实现复用 Insert 接口：

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

如果 Key 已存在，Insert 返回失败，但迭代器指向该节点，直接返回 Value 引用；如果不存在，则插入默认构造的 Value，并返回新节点的 Value 引用。

typename 的作用

在模板编程中，typename 关键字至关重要。当我们在模板类内部使用嵌套类型名称（如 RBTree<...>::Iterator）时，编译器在实例化前无法确定该名称是类型还是静态成员变量。

加上 typename 后，我们显式告诉编译器：'无论模板参数如何变化，这个名称始终是一个确定的类型'。这是编写泛型代码时的强制要求，特别是在依赖模板参数的嵌套名称上。

完整代码示例

以下是经过整合与优化的核心代码，包含了红黑树节点、迭代器及容器封装。

RBTree.h

#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), _col(RED) {}
};

template<class T, class Ref, class Ptr>
struct RBTreeIterator {
    typedef RBTreeNode<T> Node;
    typedef RBTreeIterator<T, Ref, Ptr> Self;

    Node* _root;
    Node* _node;

    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++(int) {
        Self tmp = *this;
        ++(*this);
        return tmp;
    }

    Self& operator--() {
        if (_node == nullptr) {
            Node* maxRight = _root;
            while (maxRight && maxRight->_right)
                maxRight = maxRight->_right;
            _node = maxRight;
        } 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;
    }

    Self& operator--(int) {
        Self tmp = *this;
        --(*this);
        return tmp;
    }

    bool operator==(const Self& s) const { return _node == s._node; }
    bool operator!=(const Self& s) const { 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*> ConstIterator;

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

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

    ConstIterator Begin() const {
        Node* minLeft = _root;
        while (minLeft && minLeft->_left) minLeft = minLeft->_left;
        return ConstIterator(minLeft, _root);
    }

    ConstIterator End() const { return ConstIterator(nullptr, _root); }

    pair<Iterator, bool> Insert(const T& data) {
        if (_root == nullptr) {
            _root = new Node(data);
            _root->_col = BLACK;
            return { Iterator(_root, _root), true };
        }

        Node* parent = nullptr;
        Node* cur = _root;
        KeyOft kot;
        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 { Iterator(cur, _root), 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;

        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 { Iterator(newNode, _root), true };
    }

    Iterator Find(const K& key) {
        Node* cur = _root;
        KeyOft kot;
        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();
    }

private:
    void RotateR(Node* parent) {
        Node* subL = parent->_left;
        Node* subLR = subL->_right;
        parent->_left = subLR;
        if (subLR) subLR->_parent = parent;
        Node* ppnode = parent->_parent;
        subL->_right = parent;
        parent->_parent = subL;
        if (parent == _root) {
            _root = subL;
            subL->_parent = nullptr;
        } else {
            if (ppnode->_left == parent) ppnode->_left = subL;
            else ppnode->_right = subL;
            subL->_parent = ppnode;
        }
    }

    void RotateL(Node* parent) {
        Node* subR = parent->_right;
        Node* subRL = subR->_left;
        parent->_right = subRL;
        if (subRL) subRL->_parent = parent;
        Node* ppnode = parent->_parent;
        subR->_left = parent;
        parent->_parent = subR;
        if (parent == _root) {
            _root = subR;
            subR->_parent = nullptr;
        } else {
            if (ppnode->_left == parent) ppnode->_left = subR;
            else ppnode->_right = subR;
            subR->_parent = ppnode;
        }
    }

    Node* _root = nullptr;
};

myset.h

#pragma once
#include "RBTree.h"
namespace wzy {
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>::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& key) { return _t.Insert(key); }
    iterator find(const K& key) { return _t.Find(key); }

private:
    RBTree<K, const K, SetKeyOft> _t;
};
}

mymap.h

#pragma once
#include "RBTree.h"
namespace wzy {
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>::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); }
    iterator find(const K& key) { return _t.Find(key); }

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

private:
    RBTree<K, pair<const K, V>, MapKeyOft> _t;
};
}

至此，基于红黑树的关联式容器封装已完成。通过上述实现，我们不仅掌握了红黑树的平衡算法，还深入理解了 C++ 模板元编程、迭代器模式以及容器设计的核心思想。