C++ 红黑树封装实战：从零实现 Map 与 Set

iterator 核心代码基于红黑树节点指针封装。

2.3 迭代器 iterator 实现思路分析

iterator 实现的大框架跟 list 的 iterator 思路一致，用一个类型封装结点的指针，再通过重载运算符实现迭代器有像指针一样访问的行为。

难点是 operator++ 和 operator-- 的实现。map 和 set 的迭代器走的是中序遍历（左子树 -> 根结点 -> 右子树），begin() 会返回中序第一个结点的 iterator。

迭代器 ++ 的核心逻辑就是不看全局，只看局部，只考虑当前中序局部要访问的下一个结点。

1. 如果 it 指向的结点的右子树不为空，代表当前结点已经访问完了，要访问下一个结点是右子树的中序第一个（最左结点），直接找右子树的最左结点即可。
2. 如果 it 指向的结点的右子树为空，代表当前结点已经访问完了且当前结点所在的子树也访问完了，要访问的下一个结点在当前结点的祖先里面，沿着当前结点到根的祖先路径向上找。
   • 如果当前结点是父亲的左，根据中序左子树 -> 根结点 -> 右子树，下一个访问的结点就是当前结点的父亲。
   • 如果当前结点是父亲的右，当前结点所在子树访问完了，当前结点所在父亲的子树也访问完了，那么下一个访问的需要继续往根的祖先中去找，直到找到孩子是父亲左的那个祖先。

end() 如何表示？当 it 指向最后一个结点时，++it 后找不到孩子是父亲左的祖先，父亲为空，把 it 中的结点指针置为 nullptr，用 nullptr 充当 end。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.4 map 支持 []

1. map 要支持 [] 主要需要修改 insert 返回值支持，修改 RBtree 中的 insert 返回值为 pair<Iterator, bool>。
2. map 支持 []，insert 支持 [] 实现很简单。

完整代码示例与实践演示

RBTree.h

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) { } 
    Ref operator*() { return _node->_data; } 
    Ptr operator->() { return &_node->_data; } 
    Self& operator++() { 
        if (_node->_right) { 
            Node* minRight = _node->_right; 
            while (minRight->_left) { minRight = minRight->_left; } 
            _node = minRight; 
        } else { 
            Node* cur = _node; 
            Node* parent = cur->_parent; 
            while (parent && cur == parent->_right) { 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> struct RBTree { 
    typedef RBTreeNode<T> Node; 
public: 
    typedef RBTreeIterator<T, T&, T*> Iterator; 
    typedef RBTreeIterator<T, const T&, const T*> ConstIterator; 
    ~RBTree() { Destroy(_root); _root = nullptr; } 
    Iterator Begin() { Node* minLeft = _root; while (minLeft && minLeft->_left) { minLeft = minLeft->_left; } return Iterator(minLeft); } 
    Iterator End() { return Iterator(nullptr); } 
    ConstIterator Begin() const { Node* minLeft = _root; while (minLeft && minLeft->_left) { minLeft = minLeft->_left; } return ConstIterator(minLeft); } 
    ConstIterator End() const { return ConstIterator(nullptr); } 
    pair<Iterator, bool> Insert(const T& data) { 
        if (_root == nullptr) { 
            _root = new Node(data); 
            _root->_col = BLACK; 
            return { 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(data) < kot(cur->_data)) { parent = cur; cur = cur->_left; } 
            else { return { Iterator(cur), 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 (grandfather->_left == parent) { 
                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), true }; 
    } 
    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 (parent == _root) { _root = subL; subL->_parent = nullptr; } 
        else { 
            if (parentParent->_left == parent) { parentParent->_left = subL; } 
            else { parentParent->_right = subL; } 
            subL->_parent = parentParent; 
        } 
    } 
    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 == nullptr) { _root = subR; subR->_parent = nullptr; } 
        else { 
            if (parent == parentParent->_left) { parentParent->_left = subR; } 
            else { parentParent->_right = subR; } 
            subR->_parent = parentParent; 
        } 
    } 
    Iterator Find(const K& key) { 
        KeyOfT kot; 
        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); } 
        } 
        return End(); 
    } 
    int Height() { return _Height(_root); } 
    int Size() { return _Size(_root); } 
private: 
    int _Size(Node* root) { if (root == nullptr) return 0; return _Size(root->_left) + _Size(root->_right) + 1; } 
    int _Height(Node* root) { if (root == nullptr) return 0; int leftHeight = _Height(root->_left); int rightHeight = _Height(root->_right); return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1; } 
    void Destroy(Node* root) { if (root == nullptr) return; Destroy(root->_left); Destroy(root->_right); delete root; } 
private: 
    Node* _root = nullptr; 
};

Map.h

#pragma once 
#include"RBTree.h" 
namespace jqj { 
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) { 
        auto [it, flag] = _t.Insert({ key, V() }); 
        return it->second; 
    } 
private: 
    RBTree<K, pair<const K, V>, MapKeyOfT> _t; 
}; 
}

Set.h

#pragma once 
#include"RBTree.h" 
namespace jqj { 
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; 
}; 
}

Test.cpp

#define _CRT_SECURE_NO_WARNINGS 1 
#include<iostream> 
#include<vector> 
using namespace std; 
#include"RBTree.h" 
#include"Map.h" 
#include"Set.h" 

template<class T> void func(const jqj::set<T>& s) { 
    typename jqj::set<T>::const_iterator it = s.begin(); 
    while (it != s.end()) { 
        cout << *it << " "; 
        ++it; 
    } 
    cout << endl; 
} 

void Test_set() { 
    jqj::set<int> s; 
    s.insert(1); s.insert(2); s.insert(1); s.insert(5); s.insert(0); s.insert(10); s.insert(8); 
    jqj::set<int>::iterator it = s.begin(); 
    while (it != s.end()) { 
        cout << *it << " "; 
        ++it; 
    } 
    cout << endl; 
    func(s); 
} 

void Test_map() { 
    jqj::map<string, string> dict; 
    dict.insert({ "sort", "排序" }); 
    dict.insert({ "left", "左边" }); 
    dict.insert({ "right", "右边" }); 
    dict["string"] = "字符串"; // 插入 + 修改 
    dict["left"] = "左边 xxx"; // 修改 
    auto it = dict.begin(); 
    while (it != dict.end()) { 
        it->second += 'x'; 
        cout << it->first << ":" << it->second << endl; 
        ++it; 
    } 
    cout << endl; 
    for (auto& [k, v] : dict) { 
        cout << k << ":" << v << endl; 
    } 
    cout << endl; 
    string arr[] = { "苹果", "西瓜", "苹果", "西瓜", "苹果", "苹果", "西瓜", "苹果", "香蕉", "苹果", "香蕉" }; 
    jqj::map<string, int> countMap; 
    for (auto& e : arr) { 
        countMap[e]++; 
    } 
    for (auto& [k, v] : countMap) { 
        cout << k << ":" << v << endl; 
    } 
    cout << endl; 
} 

int main() { 
    Test_set(); 
    Test_map(); 
    return 0; 
}

运行结果

测试用例验证了 set 的去重有序输出以及 map 的键值对插入、修改和统计功能。