C++ STL map与set底层结构及迭代器实现详解

1.map、set的基本结构

通过查看官方文档，截取部分关键代码，我们可以发现 set 虽然是 K-K 类型，map 是 K-V 类型，但是实际上这两个类共用一个红黑树，准确来说是共用同一个模板类型，set 是 <K, K>，map 是 <K, pair<K, V>>，下面会进行详细解析。

• size_type node_count：用于记录红黑树节点数量，跟踪树的大小
• link_type header：是指向红黑树头节点的指针
• Value value_field：存储节点的值

那么下面我们将自己实现简单的 set 和 map 类：

2.map、set模拟实现

2.1 初步定义

template<class K> class set {
private: 
    RBTree<K, K> _t;
};

template<class K, class V> class map {
private: 
    RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};

平常我们认为键值对指的就是 K 和 V，但是在库里不是这样的，库里的 K 表示键值对的类型，V 表示插入红黑树的键值对，只不过对于 set 来说，K 和 V 是一样的。

在红黑树中，定义的模板参数 T，而不是原先的 pair，这里的 T 表示插入的数据 _data 的类型，这种定义方法能够共同使用同一参数模板，避免额外的代码编写。

2.2 仿函数实现

template<class K> class set {
    struct SetKeyOfT {
        const K& operator()(const K& key) {
            return key;
        }
    };
private: 
    RBTree<K, K, SetKeyOfT> _t;
};

template<class K, class V> class map {
    struct MapKeyOfT {
        const K& operator()(const pair<K, V>& kv) {
            return kv.first;
        }
    };
private: 
    RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};

我们知道 set 和 map 是通过比较 key，在红黑树中来插入的，但是由于上述的定义，如果每次对于 map 都频繁取出 first 就太麻烦了，因此就定义了仿函数。

为什么使用仿函数而不是普通函数呢？

红黑树中只要涉及到数据 _data 的地方，就需要使用到仿函数提取 key，使用普通函数消耗太大，而仿函数带有 inline 的性质，降低消耗。同时官方文档中还对比较进行了实现，即 Compare，模板要求参数必须是一个类型，而普通函数无法作为类型传递。

为什么要自己定义仿函数，pair自带的仿函数不行吗？

虽然 pair 确实有自己的仿函数比较，但是他是比较完 first 后不行，会接着比较 second，这不符合我们的设计思路。

截取了部分 insert 中的代码，利用仿函数确实是能够简单的实现键值 first 的提取，我们再对整体的调用思路进行整理。

其实仿函数主要是为了 map 而设计的，为的就是提取 first，set 为了保持设计模式的一致，因而也设计了相同的仿函数，这样就不用关心是否需要调用这一点了，保持一致性。

这里我们不对 Compare 进行实现，有兴趣的可以自己去看底层代码。

值得注意的是： 仿函数内不实现比较功能是因为，比较功能是一个外层调用功能，如果放在内部就不能操作者自行去调用了，况且 Compare 也是以仿函数的形式实现的，两个仿函数嵌套过于复杂，不好使用。

2.3 Find功能实现

Node* 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 cur;
        }
    }
    return nullptr;
}

2.4 迭代器初步功能实现

类似的迭代器分析我们在 list 部分有做过解析，确实大体上是相像的，但是结构并不一样，这里的树形结构需要以中序遍历：左 - 根 - 右的方式遍历。

template<class T> struct __TreeIterator {
    typedef RBTreeNode<T> Node;
    typedef __TreeIterator<T> Self;
    Node* _node;
    __TreeIterator(Node* node):_node(node){}
};

库里的迭代器模式并不能满足我们的设计需要，所以这里自己构建一个 __TreeIterator 类。

2.4.1 ++运算符重载

Self& operator++() {
    if(_node->_right) {
        // 右树的最左节点 (最小节点)
        Node* subLeft = _node->_right;
        while(subLeft->_left) {
            subLeft = subLeft->_left;
        }
        _node = subLeft;
    } else {
        Node* cur = _node;
        Node* parent = cur->_parent;
        // 找孩子是父亲左的那个祖先节点，就是下一个要访问的节点
        while(parent) {
            if(cur == parent->_left) {
                break;
            } else {
                cur = cur->_parent; 
                parent = parent->_parent;
            }
        }
        _node = parent;
    }
    return *this;
}

中序遍历的方式是 左 - 根 - 右，因此可以总结为两种情况来遍历：

• 当前节点有右子树
  • 处理方式： 找到右子树的最左节点（即右子树中的最小值）
  • 原因： 在中序遍历中，当前节点的下一个节点是其右子树的最左节点
• 当前节点没有右子树
  • 处理方式： 向上回溯，直到找到某个祖先节点，使得当前节点位于该祖先的左子树中
  • 原因： 在中序遍历中，若无右子树，则下一个节点是第一个满足'当前节点是其左子节点'的祖先

值得注意的是： 当前节点没有右子树的情况，是 左 - 根 - 右 的最后一步，无论是在根的左边还是右边，最终都会回到根节点，所以直接 _node = parent 即可。

2.4.2 --运算符重载

Self& operator--() {
    if(_node->_left) {
        Node* subRight = _node->_left;
        while(subRight->_right) {
            subRight = subRight->_right;
        }
        _node = subRight;
    } else {
        // 孩子是父亲的右的那个节点
        Node* cur = _node;
        Node* parent = cur->_parent;
        while(parent && cur == parent->_left) {
            cur = cur->_parent; 
            parent = parent->_parent;
        }
        _node = parent;
    }
    return *this;
}

operator-- 的思路和 operator++ 是一样的，反过来遍历就行了。

2.4.3 *运算符重载

T& operator*() {
    return _node->_data;
}

2.4.4 ->运算符重载

T* operator->() {
    return &_node->_data;
}

这里再提醒一下重载 -> 是因为用 * 的代码不够简洁，具体分析参考 list 部分的解析。

2.4.5 !=运算符重载

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

_node：当前迭代器指向的节点 s._node：另一个迭代器（作为参数传入）指向的节点

2.4.6 begin()

//RBTree.h
iterator begin() {
    Node* leftMin = _root;
    while(leftMin && leftMin->_left) {
        leftMin = leftMin->_left;
    }
    return iterator(leftMin);
}
//Set.h / Map.h
iterator begin() {
    return _t.begin();
}

2.4.7 end()

//RBTree.h
iterator end() {
    return iterator(nullptr);
}
//Set.h / Map.h
iterator end() {
    return _t.end();
}

现在已经可以基本实现遍历的功能了。

2.5 迭代器进阶功能实现

2.5.1 set：const迭代器及insert的实现

typedef typename RBTree<K, K, SetKeyOfT>::const_iterator iterator;
typedef typename RBTree<K, K, SetKeyOfT>::const_iterator const_iterator;
const_iterator begin() const {
    return _t.begin();
}
const_iterator end() const {
    return _t.end();
}

由于 set 规定 key 是不可以被修改的，因此 iterator 和 const_iterator 本质上其实都是 const_iterator。

值得注意的是： begin() 和 end() 的 const 迭代器函数被 const 修饰是为了满足常量容器对象或非常量容器对象都能调用。

insert 的错误代码：

pair<iterator, bool> insert(const K& key) {
    return _t.Insert(key);
}

这里是返回红黑树的插入，红黑树的插入详见下面的代码展示。

从之前的学习我们知道 insert 返回的是 pair<iterator, bool>，那么是不是直接返回 insert 的结果就好了呢？看似确实是没问题，但是这里理了个巨大的坑，我们实际分析一波：

• _t.Insert(key) 返回的是 RBTree::iterator，是一个普通迭代器
• pair<iterator, bool> insert(const K& key) 返回的是 set::iterator，是一个 const 迭代器

insert 的正确代码：

// iterator RBTree::const_iterator
pair<iterator, bool> insert(const K& key) {
    // pair<RBTree::iterator, bool>
    pair<typename RBTree<K, K, SetKeyOfT>::iterator, bool> ret = _t.Insert(key);
    return pair<iterator, bool>(ret.first, ret.second);
}

正确的做法是先将 insert 返回的普通迭代器由变量 ret 存储，然后再用一个匿名对象进行构造，将 ret 的普通迭代器构造成 const 迭代器返回即可，下面将进行详细的构造原理解释：

回看官方文档发现 iterator 和 const_iterator 都是被单独拿出来实例化的，并没有受到 Ref 和 Ptr 的影响，那么此时就分为两种情况：

• 普通迭代器的拷贝构造
  • 当 __rb_tree_iterator 是普通迭代器时，iterator 就是自身类型，此时构造函数等价于：

    __rb_tree_iterator(const __rb_tree_iterator<Value, Value&, Value*>& it):node(it.node){}
    

    这是一个标准的拷贝构造函数，用于创建一个新的普通迭代器，指向相同的节点
• const迭代器的构造
  • 当 __rb_tree_iterator 是 const 迭代器时，iterator 指的是普通迭代器类型，此时构造函数等价于：

    __rb_tree_iterator(const __rb_tree_iterator<Value, Value&, Value*>& it):node(it.node){}
    

    这变成了一个构造函数，允许从普通迭代器创建 const 迭代器

所以可以理解为单独拿出来实例化是为了不让 Ref 和 Ptr 影响参数，而外面的类型就会受 Ref 和 Ptr 影响，这样就能保证外面的类型是 const 迭代器，里面的参数是普通迭代器，成功构造出一个支持普通迭代器构造 const 迭代器的构造函数。

那再转到实际代码上，ret.first 的类型是 typename RBTree<K, K, SetKeyOfT>::iterator，返回值 pair 的第一个元素类型是 set 类中定义的 iterator，实际上是 typename RBTree<K, K, SetKeyOfT>::const_iterator。

ret.first 会调用自定义的迭代器类型的构造函数 __TreeIterator(const Iterator& it) 进行单参数转换，变成 const_iterator。

2.5.2 map：const迭代器及insert、[ ]运算符重载的实现

typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::iterator iterator;
typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::const_iterator 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();
}

对于 map 来说，key 是不允许改变的，value 是可以改变的，但是如果像 set 那样写的话 key 和 value 都不能修改了，所以直接在 pair 的 key 加 const，控制 value 即可。

insert 代码：

pair<iterator, bool> insert(const pair<K, V>& kv) {
    return _t.Insert(kv);
}

map 就没有像 set 那么麻烦了，红黑树和 map 的迭代器是一致的。

[ ]运算符重载 代码：

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

之前详细解释过，可以看之前的博客。

3.代码展示

MySet.h

#pragma once
#include"RBTree.h"
namespace bit {
template<class K> class set {
    struct SetKeyOfT {
        const K& operator()(const K& key) {
            return key;
        }
    };
public:
    typedef typename RBTree<K, K, SetKeyOfT>::const_iterator iterator;
    typedef typename RBTree<K, K, SetKeyOfT>::const_iterator const_iterator;
    const_iterator begin() const {
        return _t.begin();
    }
    const_iterator end() const {
        return _t.end();
    }
    // iterator RBTree::const_iterator
    pair<iterator, bool> insert(const K& key) {
        // pair<RBTree::iterator, bool>
        pair<typename RBTree<K, K, SetKeyOfT>::iterator, bool> ret = _t.Insert(key);
        return pair<iterator, bool>(ret.first, ret.second);
    }
private: 
    RBTree<K, K, SetKeyOfT> _t;
};
}

MyMap.h

#pragma once
#include"RBTree.h"
namespace bit {
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>::const_iterator 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();
    }
    V& operator[](const K& key) {
        pair<iterator, bool> ret = insert(make_pair(key, V()));
        return ret.first->second;
    }
    pair<iterator, bool> insert(const pair<K, V>& kv) {
        return _t.Insert(kv);
    }
private: 
    RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};
}

RBTree.h

#pragma once
#include<iostream>
using namespace std;
enum Colour { RED, BLACK };

template<class T> struct RBTreeNode {
    RBTreeNode<T>* _left;
    RBTreeNode<T>* _right;
    RBTreeNode<T>* _parent;
    T _data;
    Colour _col;
    RBTreeNode(const T& data):_left(nullptr),_right(nullptr),_parent(nullptr),_data(data),_col(RED){}
};

template<class T, class Ptr, class Ref> struct __TreeIterator {
    typedef RBTreeNode<T> Node;
    typedef __TreeIterator<T, Ptr, Ref> Self;
    typedef __TreeIterator<T, T*, T&> Iterator;
    __TreeIterator(const Iterator& it):_node(it._node){}
    Node* _node;
    __TreeIterator(Node* node):_node(node){}
    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;
    }
    Self& operator--() {
        if(_node->_left) {
            Node* subRight = _node->_left;
            while(subRight->_right) {
                subRight = subRight->_right;
            }
            _node = subRight;
        } else {
            // 孩子是父亲的右的那个节点
            Node* cur = _node;
            Node* parent = cur->_parent;
            while(parent && cur == parent->_left) {
                cur = cur->_parent; 
                parent = parent->_parent;
            }
            _node = parent;
        }
        return *this;
    }
    Self& operator++() {
        if(_node->_right) {
            // 右树的最左节点 (最小节点)
            Node* subLeft = _node->_right;
            while(subLeft->_left) {
                subLeft = subLeft->_left;
            }
            _node = subLeft;
        } else {
            Node* cur = _node;
            Node* parent = cur->_parent;
            // 找孩子是父亲左的那个祖先节点，就是下一个要访问的节点
            while(parent && cur == parent->_right) {
                cur = cur->_parent; 
                parent = parent->_parent;
            }
            _node = parent;
        }
        return *this;
    }
};
// set->RBTree<K, K, SetKeyOfT> _t;
// map->RBTree<K, pair<K, V>, MapKeyOfT> _t;

template<class K, class T, class KeyOfT> struct RBTree {
    typedef RBTreeNode<T> Node;
public:
    // 同一个类模板，传的不同的参数实例化出的不同类型
    typedef __TreeIterator<T, T*, T&> iterator;
    typedef __TreeIterator<T, const T*, const T&> const_iterator;
    iterator begin() {
        Node* leftMin = _root;
        while(leftMin && leftMin->_left) {
            leftMin = leftMin->_left;
        }
        return iterator(leftMin);
    }
    iterator end() {
        return iterator(nullptr);
    }
    const_iterator begin() const {
        Node* leftMin = _root;
        while(leftMin && leftMin->_left) {
            leftMin = leftMin->_left;
        }
        return const_iterator(leftMin);
    }
    const_iterator end() const {
        return const_iterator(nullptr);
    }
    Node* 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 cur;
            }
        }
        return nullptr;
    }
    pair<iterator, bool> Insert(const T& data) {
        if(_root == nullptr) {
            _root = new Node(data);
            _root->_col = BLACK;
            return make_pair(iterator(_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 make_pair(iterator(cur), false);
            }
        }
        cur = new Node(data);
        cur->_col = RED;
        Node* newnode = cur;
        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;
                // u存在且为红
                if(uncle && uncle->_col == RED) {
                    // 变色
                    parent->_col = uncle->_col = BLACK;
                    grandfather->_col = RED;
                    // 继续向上处理
                    cur = grandfather;
                    parent = cur->_parent;
                } else // u不存在 或 存在且为黑
                {
                    if(cur == parent->_left) {
                        // g
                        // p
                        // c
                        RotateR(grandfather);
                        parent->_col = BLACK;
                        grandfather->_col = RED;
                    } else {
                        // g
                        // p
                        // c
                        RotateL(parent);
                        RotateR(grandfather);
                        cur->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    break;
                }
            } else // parent == grandfather->_right
            {
                Node* uncle = grandfather->_left;
                // u存在且为红
                if(uncle && uncle->_col == RED) {
                    // 变色
                    parent->_col = uncle->_col = BLACK;
                    grandfather->_col = RED;
                    // 继续向上处理
                    cur = grandfather;
                    parent = cur->_parent;
                } else {
                    if(cur == parent->_right) {
                        // g
                        // p
                        // c
                        RotateL(grandfather);
                        grandfather->_col = RED;
                        parent->_col = BLACK;
                    } else {
                        // g
                        // p
                        // c
                        RotateR(parent);
                        RotateL(grandfather);
                        cur->_col = BLACK;
                        grandfather->_col = RED;
                    }
                    break;
                }
            }
        }
        _root->_col = BLACK;
        return make_pair(iterator(newnode), true);
    }
    void RotateL(Node* parent) {
        ++_rotateCount;
        Node* cur = parent->_right;
        Node* curleft = cur->_left;
        parent->_right = curleft;
        if(curleft) {
            curleft->_parent = parent;
        }
        cur->_left = parent;
        Node* ppnode = parent->_parent;
        parent->_parent = cur;
        if(parent == _root) {
            _root = cur;
            cur->_parent = nullptr;
        } else {
            if(ppnode->_left == parent) {
                ppnode->_left = cur;
            } else {
                ppnode->_right = cur;
            }
            cur->_parent = ppnode;
        }
    }
    void RotateR(Node* parent) {
        ++_rotateCount;
        Node* cur = parent->_left;
        Node* curright = cur->_right;
        parent->_left = curright;
        if(curright) curright->_parent = parent;
        Node* ppnode = parent->_parent;
        cur->_right = parent;
        parent->_parent = cur;
        if(ppnode == nullptr) {
            _root = cur;
            cur->_parent = nullptr;
        } else {
            if(ppnode->_left == parent) {
                ppnode->_left = cur;
            } else {
                ppnode->_right = cur;
            }
            cur->_parent = ppnode;
        }
    }
    // 17:20继续
    bool CheckColour(Node* root, int blacknum, int benchmark) {
        if(root == nullptr) {
            if(blacknum != benchmark) return false;
            return true;
        }
        if(root->_col == BLACK) {
            ++blacknum;
        }
        if(root->_col == RED && root->_parent && root->_parent->_col == RED) {
            cout << root->_kv.first << "出现连续红色节点" << endl;
            return false;
        }
        return CheckColour(root->_left, blacknum, benchmark) && CheckColour(root->_right, blacknum, benchmark);
    }
    bool IsBalance() {
        return IsBalance(_root);
    }
    bool IsBalance(Node* root) {
        if(root == nullptr) return true;
        if(root->_col != BLACK) {
            return false;
        }
        // 基准值
        int benchmark = 0;
        Node* cur = _root;
        while(cur) {
            if(cur->_col == BLACK) ++benchmark;
            cur = cur->_left;
        }
        return CheckColour(root, 0, benchmark);
    }
    int Height() {
        return Height(_root);
    }
    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;
    }
private:
    Node* _root = nullptr;
public:
    int _rotateCount = 0;
};