1 封装红黑树实现 set 和 map
1.1 对底层源码及框架分析
SGI-STL30 版本源代码中，map 和 set 的源代码在 stl_map.h/stl_set.h/stl_tree.h 等头文件中。核心部分如下：
// set
#ifndef __SGI_STL_INTERNAL_TREE_H
#include <stl_tree.h>
#endif
#include <stl_set.h>
#include <stl_multiset.h>
// map
#ifndef __SGI_STL_INTERNAL_TREE_H
#include <stl_tree.h>
#endif
#include <stl_map.h>
#include <stl_multimap.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:
typedef rb_tree<key_type, value_type, identity<value_type>, key_compare, Alloc> rep_type;
rep_type t; // red-black tree representing set
};
// 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:
typedef rb_tree<key_type, value_type, select1st<value_type>, key_compare, Alloc> rep_type;
rep_type t; // red-black tree representing map
};
// stl_tree.h
struct __rb_tree_node_base {
typedef __rb_tree_color_type color_type;
typedef __rb_tree_node_base* base_ptr;
color_type color;
base_ptr parent;
base_ptr left;
base_ptr right;
};
template<class Key, class Value, class KeyOfValue, class Compare, class Alloc = alloc>
class rb_tree {
protected:
typedef void* void_pointer;
typedef __rb_tree_node_base* base_ptr;
typedef __rb_tree_node<Value> rb_tree_node;
typedef rb_tree_node* link_type;
typedef Key key_type;
typedef Value value_type;
public:
pair<iterator, bool> insert_unique(const value_type& x);
size_type erase(const key_type& x);
iterator find(const key_type& x);
protected:
size_type node_count;
link_type header;
};
template<class Value>
struct __rb_tree_node : public __rb_tree_node_base {
typedef __rb_tree_node<Value>* link_type;
Value value_field;
};
- 通过框架分析，源码中 rb_tree 用了巧妙的泛型思想实现。
- set 实例化 rb_tree 时第二个模板参数给的是 key，map 实例化时给的是 pair<const key, T>。
- 源码中的 value_type 是红黑树结点中存储的真实数据类型。
- 第一个模板参数 Key 用于 find/erase 等函数做形参类型。
2 模拟实现 map 和 set
2.1 实现出复用红黑树的框架，并支持 insert
参考源码框架，map 和 set 复用之前实现的红黑树。相比源码调整一下，key 参数用 K，value 参数用 V，红黑树中的数据类型使用 T。因为 RBTree 实现了泛型不知道 T 参数导致是 K 还是 pair<K, V>，insert 内部进行比较时无法直接比较，需要我们在 map 和 set 层分别实现一个 MapKeyofpair 和 SetKeyofpair 的仿函数传给 RBtree 的 KeyOfPair。
// map.h
template<class T1, class T2>
bool operator<(const pair<T1, T2>& lhs, const pair<T1, T2>& rhs) {
return lhs.first < rhs.first || (!(rhs.first < lhs.first) && lhs.second < rhs.second);
}
template<class K, class V>
class map {
struct MapKeyofpair {
const K& operator()(const pair<K, V>& kv) { return kv.first; }
};
public:
bool insert(const pair<K, V>& kv) {
return _t.Insert(kv); // 底层是红黑树，调用红黑树的 insert 就行了
}
private:
RBTree<K, pair<K, V>, MapKeyofpair> _t;
};
// set.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, SetKeyofpair> _t;
};
// RBtree.h
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) {}
};
template<class K, class T, class KeyOfT>
class RBtree {
private:
typedef RBTreeNode<T> Node;
Node* _root = nullptr;
public:
bool Insert(const T& data) {
if (_root == nullptr) {
_root = new Node(data);
_root->_col = BLACK;
return 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 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;
// ... 平衡处理逻辑省略
return true;
}
};
2.2 支持 iterator 的实现
红黑树的迭代器不是 ++ 就能得到下一个结点的。Iterator 实现的框架跟 list 的 iterator 思路一致，用一个类型封装结点的指针再通过重载运算符实现。
2.2.1 红黑树迭代器结构
template<class T, class Ref, class Ptr>
struct RBtreeIterator {
typedef RBTreeNode<T> Node;
typedef RBtreeIterator<T, Ref, Ptr> Self;
Node* _node;
Node* _root;
RBtreeIterator(Node* node, Node* root) : _node(node), _root(root) {}
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; }
};
2.2.2 迭代器++
迭代器++的核心逻辑只看局部，考虑当前中序局部要访问的下一个结点（左子树 -> 根结点 -> 右子树）。
- 情况一（右子树不为空）：找右子树的最左结点。
- 情况二（右子树为空）：沿着祖先路径向上找，直到找到孩子是父亲左的那个祖先。
Self operator++() {
if (_node->_right) {
Node* min = _node->_right;
while (min->_left) { min = min->_left; }
_node = min;
} else {
Node* cur = _node;
Node* parent = cur->_parent;
while (parent && cur == parent->_right) {
cur = parent;
parent = cur->_parent;
}
_node = parent;
}
return *this;
}
end() 实现：当 it 指向 50 时，++it 找不到孩子是父亲左的祖先，父指针为空，将 it 中的结点指针置为 nullptr 充当 end。
2.2.4 iterator--
迭代器--的实现跟++的思路完全类似，逻辑正好反过来（右子树 -> 根结点 -> 左子树）。
Self operator--() {
if (_node == nullptr) {
Node* rightMost = _root;
while (rightMost && rightMost->_right) { rightMost = rightMost->_right; }
_node = rightMost;
} 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;
}
3 注意须知 [实现 map/set]
- 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;
3.1 map[] 实现
map 要支持 [] 主要需要修改 insert 返回值支持：
V& operator[](const K& key) {
pair<iterator, bool> ret = insert({ key, V() });
return ret.first->second;
}
3.2 代码实现
// map.h
template<class K, class V>
class map {
struct mapKeyofpair {
const K& operator()(const pair<K, V>& kv) { return kv.first; }
};
public:
typedef typename RBtree<K, pair<const K, V>, mapKeyofpair>::Iterator iterator;
typedef typename RBtree<K, pair<const K, V>, mapKeyofpair>::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); }
private:
RBtree<K, pair<const K, V>, mapKeyofpair> _t;
};
// set.h
template<class K>
class set {
struct setKeyofpair {
const K& operator()(const K& key) { return key; }
};
public:
typedef typename RBtree<K, const K, setKeyofpair>::Iterator iterator;
typedef typename RBtree<K, const K, setKeyofpair>::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& data) { return _t.insert(data); }
private:
RBtree<K, const K, setKeyofpair> _t;
};
// RBtree.h 核心逻辑包含旋转、变色、插入等完整实现...
