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

源码中的 rb_tree 类模板实现了非常巧妙的泛型思想。它不直接写死存储类型，而是由第二个模板参数 value 决定 _rb_tree_node 中存储的数据类型。

• set 实例化 rb_tree：第二个模板参数给的是 Key。
• map 实例化 rb_tree：第二个模板参数给的是 pair<const key, T>。

这样一棵红黑树既能实现 Key 搜索场景的 set，也能实现 Key/Value 搜索场景的 map。

注意，源码中模板参数用了 T 代表 value，而内部写的 value_type 并非日常 Key/Value 场景中的 value，而是红黑树节点中存储的真实数据类型。

rb_tree 第二个模板参数 value 控制了节点存储的数据类型，为什么还要穿第一个模板参数 Key？尤其是 set，两个参数是一样的。

这是因为对于 map 和 set，find / erase 时的模板参数都是 Key。所以第一个模板参数是传给这些函数做形参类型的。对于 set 而言，两个参数是一样的；但对于 map 而言就不一样了，map 容器 insert 的是 pair 对象，但是 find / erase 的是 Key 对象。

2. map 和 set 的模拟实现

2.1 实现出复用红黑树的框架（支持 insert）

参考前面的源码框架，map 和 set 确实是复用之前实现的红黑树。这里相比源码调整一下，key 参数用 K，value 参数用 V，红黑树中的数据类型使用 T。

因为 RBTree 实现了泛型，不知道 T 参数导致是 K，还是 pair<K, V>，insert 内部进行插入逻辑比较时，就没办法进行比较。因为 pair 的默认支持的是 key 和 value 一起参与比较，我们需要的是任何时候只比较 key。

所以我们在 map 和 set 层分别实现一个 MapKeyOfT 和 SetKeyOfT 的仿函数传给 RBTree 的 KeyOfT，然后 RBTree 中通过 KeyOfT 仿函数取出 T 类型对象中的 key，再进行比较。

2.2 迭代器 iterator 的实现

iterator 核心代码——

2.3 迭代器 iterator 实现思路分析

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

这里的难点是 operator++ 和 operator-- 的实现。在使用部分，我们分析 map 和 set 的迭代器走的是中序遍历：左子树 -> 根结点 -> 右子树。那么 begin() 会返回中序第一个结点的 iterator，也就是最左结点的迭代器。

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

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

end() 如何表示呢？ 当 it 指向 50 时，++it 时，50 是 40 的右，40 是 30 的右，30 是 18 的右，18 到根没有父亲，没有找到孩子是父亲左的那个祖先，这时父亲为空了。我们就把 it 中的结点指针置为 nullptr，我们用 nullptr 去充当 end。

需要注意的是 STL 源码中空红黑树增加了一个哨兵位头结点作为 end()，这哨兵位头结点和根互为父亲，左指向最左结点，右指向最右结点。相比我们用 nullptr 作为 end()，差别不大，他能实现的，我们也能实现。只是 --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;

如下图所示，header 是'哨兵位'，在学习数据结构中的二叉树时，我们就已经了解过哨兵位了：

2.4 map 支持 []

1. map 要支持 [] 主要需要修改 insert 返回值支持，修改 RBtree 中的 insert 返回值为：

   pair<Iterator, bool> Insert(const T& data)
   

2. map 支持 []，insert 支持 [] 实现就很简单了。

bit::map 和 bit::set 完整代码示例与实践演示

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) {
                // g
                // p u
                Node* uncle = grandfather->_right;
                // 叔叔存在且为空
                if (uncle && uncle->_col == RED) {
                    // 变色 + 继续往上处理
                    parent->_col = uncle->_col = BLACK;
                    grandfather->_col = RED;
                    cur = grandfather;
                    parent = cur->_parent;
                } else // 叔叔不存在或者叔叔存在且为黑
                {
                    // g
                    // p u
                    // c
                    // 单旋 + 变色
                    if (cur == parent->_left) {
                        RotateR(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    } else {
                        // g
                        // p u
                        // c
                        // 双旋 + 变色
                        RotateL(parent);
                        RotateR(grandfather);
                        cur->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    break;
                }
            } else {
                // g
                // u p
                Node* uncle = grandfather->_left;
                // 叔叔存在且为红，-》变色即可
                if (uncle && uncle->_col == RED) {
                    parent->_col = uncle->_col = BLACK;
                    grandfather->_col = RED;
                    // 继续往上处理
                    cur = grandfather;
                    parent = cur->_parent;
                } else {
                    // 情况二：叔叔不存在或者存在且为黑
                    // 旋转 + 变色
                    // g
                    // u p
                    // c
                    if (cur == parent->_right) {
                        RotateL(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    } else {
                        // g
                        // u p
                        // c
                        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) {
        // pair<iterator, bool> ret = _t.Insert({ key, V() });
        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()) {
        //*it = 1;
        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();
    // *it += 10;
    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->first += 'x'; // 不能修改
        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) {
        /*auto it = countMap.find(e);
        if (it != countMap.end()) {
            it->second++;
        } else {
            countMap.insert({ e, 1 });
        }*/
        countMap[e]++;
    }
    for (auto& [k, v] : countMap) {
        cout << k << ":" << v << endl;
    }
    cout << endl;
}

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

运行结果