文章目录
- 上文链接
- 一、改造红黑树
-
- [1. 分析源码](#1. 分析源码)
- [2. 整体代码框架](#2. 整体代码框架)
- [二、KeyOfT 修改 insert 中的比较逻辑](#二、KeyOfT 修改 insert 中的比较逻辑)
- [三、迭代器 iterator](#三、迭代器 iterator)
-
- [1. 整体代码框架](#1. 整体代码框架)
- [2. begin()](#2. begin())
- [3. end()](#3. end())
- [4. operator*](#4. operator*)
- [5. operator->](#5. operator->)
- [6. operator==](#6. operator==)
- [7. operator!=](#7. operator!=)
- [8. operator++](#8. operator++)
- [9. operator--](#9. operator--)
- 四、const_iterator
- [五、实现 key 不支持修改](#五、实现 key 不支持修改)
- [六、operator[ ]](#六、operator[ ])
- 七、完整代码
-
- [1. RBTree.h](#1. RBTree.h)
- [2. myset.h](#2. myset.h)
- [3. mymap.h](#3. mymap.h)
- 下文链接
上文链接
一、改造红黑树
1. 分析源码
set 与 map 的底层都是红黑树，但是我们不能直接使用一棵普通的红黑树套进去，因为 set 和 map 中所存储的数据类型是不一样的，set 中是单个值 key，而 map 中是一个 pair 类型。那么我们应该如何解决呢？我们来参考一下 STL 库中的写法。
SGI-STL30 版本源代码，set 和 map 的源代码在 map/set/stl_map.h/stl_set.h/stl_tree.h 等几个头文件 中。set 和 map 的实现结构框架核心部分截取出来如下：
cpp
// 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:
// typedefs:
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:
// typedefs:
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;
};
// stl_tree.h
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:
// insert用的是第二个模板参数做形参
pair<iterator, bool> insert_unique(const value_type& x);
// erase和find用第一个模板参数做形参
size_type erase(const key_type& x);
iterator find(const key_type& x);
protected:
size_type node_count; // keeps track of size of tree
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 用了一个巧妙的泛型思想实现，rb_tree 是实现 key 的搜索场景，还是 key/value 的搜索场景不是直接写死的，而是由第二个模板参数 Value 决定 _rb_tree_node 中存储的数据类型。
set 实例化 rb_tree 时第二个模板参数给的是 key，map 实例化 rb_tree 时第二个模板参数给的是 pair，这样一颗红黑树既可以实现key 搜索场景的 set，也可以实现 key/value 搜索场景的 map。
rb_tree 第二个模板参数 Value 已经控制了红黑树节点中存储的数据类型，为什么还要传第一个模板参数 Key 呢？尤其是 set，两个模板参数是一样的。要注意的是对于 map 和 set，find/erase 时的函数参数都是 Key，所以第一个模板参数是传给 find/erase 等函数做形参的类型的。对于 set 而言两个参数是一样的，但是对于 map 而言就完全不一样了，map 中 insert 的是 pair 对象。
2. 整体代码框架
cpp
#pragma once
#include<utility>
#include<iostream>
using namespace std;
// 枚举值表示颜色
enum Colour
{
RED,
BLACK
};
// 由于 set 和 map 底层存储的数据不一样, 可能是 key 也可能是 pair (key-value)
// 所以红黑树不是写死的, 而是写成一个模板参数 T
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)
{}
};
// 这里的 K 是 key, 提供给一些需要的函数接口
template<class K, class T>
class RBTree
{
typedef RBTreeNode<T> Node;
public:
// ...
private:
Node* _root = nullptr;
};
cpp
// myset.h
#include"RBTree.h"
namespace mine
{
template<class K>
class set
{
private:
RBTree<K, K> _t;
};
}
cpp
// mymap.h
#include"RBTree.h"
namespace mine
{
template<class K, class V>
class map
{
private:
RBTree<K, pair<K, V>> _t;
};
}二、KeyOfT 修改 insert 中的比较逻辑
由于 RBTree 实现了泛型所以我们不知道 T 参数到底是 K，还是 pair，那么 insert 内部进行比较逻辑时，就没办法按我们期望的方式进行比较，因为 pair 的默认支持的是 key 和 value 一起参与比较，而我们需要时的任何时候只比较 key。
为此，我们在 map 和 set 层分别实现一个 MapKeyOfT 和 SetKeyOfT 的仿函数传给 RBTree 的 KeyOfT，作用就是通过 KeyOfT 仿函数取出 T 类型对象中的 key，再进行比较，具体细节参考如下代码实现。
cpp
template<class K, class T, class KeyOfT>
class RBTree
{
// ...
};
cpp
namespace mine
{
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, SetKeyOfT> _t;
};
}
cpp
namespace mine
{
template<class K, class V>
class map
{
struct MapKeyOfT
{
const K& operator()(const pair<K, V>& kv)
{
return kv.first;
}
};
public:
bool Insert(const pair<K, V>& kv)
{
return _t.Insert(kv);
}
private:
RBTree<K, pair<K, V>, MapKeyOfT> _t;
};
}
cpp
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)
{
// 用 kot 取出 T 中的 key 进行比较
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);
cur->_col = RED;
if (kot(parent->_data) < kot(data)) parent->_right = cur;
else parent->_left = cur;
cur->_parent = parent;
// ...
}三、迭代器 iterator
1. 整体代码框架
iterator 实现的大框架跟 list 的 iterator 思路是一致的，用一个类封装节点的指针，再通过重载运算符实现迭代器像指针一样访问的行为。
cpp
template<class T>
struct RBTreeIterator
{
typedef RBTreeNode<T> Node;
typedef RBTreeIterator<T> Self;
Node* _node;
Node* _root; // operator-- 会用到
RBTreeIterator(Node* node, Node* root)
:_node(node)
, _root(root)
{}
// ...
};2. begin()
获取中序遍历的第一个节点的迭代器，我们只需要找到红黑树最左边的节点返回其迭代器即可。
cpp
template<class K, class T, class KeyOfT>
class RBTree
{
typedef RBTreeNode<T> Node;
public:
typedef RBTreeIterator<T> Iterator;
Iterator begin()
{
Node* minLeft = _root;
while (minLeft && minLeft->_left)
{
minLeft = minLeft->_left;
}
return Iterator(minLeft, _root);
}
// ...
};3. end()
这里我们将 nullptr 作为 end()，具体原因在 operator++ 处会提到。
cpp
Iterator end()
{
return Iterator(nullptr, _root);
}4. operator*
cpp
T& operator*()
{
return _node->_data;
}5. operator->
cpp
T* operator->()
{
return &_node->_data;
}6. operator==
cpp
bool operator==(const Self& s) const
{
return _node == s.node;
}7. operator!=
cpp
bool operator!=(const Self& s) const
{
return _node != s._node;
}8. operator++
迭代器 ++ 时，如果迭代器指向的节点的右子树不为空，代表当前节点已经访问完了，现在要访问下一个节点是右子树的中序第一个，也就是这棵树的最左节点，所以直接找右子树的最左节点即可。
如果迭代器指向的节点的右子树空，代表当前节点已经访问完了且当前节点所在的子树也访问完了，要访问的下一个节点在当前节点的祖先里面，所以要沿着当前节点到根的祖先路径向上找。那么找祖先里面的谁呢？此时相当于我们已经访问完了左根右中的左，现在要去找左根右中的根。那么就需要依次向上去访问当前节点的祖先，直到找到某个节点的父亲的左孩子是该节点，那么该节点的父亲就是我们要找的那个祖先。
cpp
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;
// parent 为空, 则 cur 是根, 说明当前树走完了, 于是把 nullptr 给 _node, nullptr 作 end()
while (parent && cur == parent->_right)
{
cur = parent;
parent = parent->_parent;
}
_node = parent;
}
return *this;
} 9. operator--
-- 和 ++ 中的逻辑是完全一样的，这里就不过多解释了，唯一需要注意的是 --end() 的情况，因为我们这里的 end() 是一个空指针，所以 --end() 需要返回的是中序遍历的最后一个节点，也就是红黑树的最右边的节点。
cpp
Self& operator--()
{
if (_node == nullptr) // end()
{
// --end()，特殊处理，走到中序最后一个节点，整棵树的最右节点
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;
}在 STL 库中，设置了一个 header 作为 end()，即一个哨兵位的节点，通过它能够去直接找到红黑树中的最左节点与最右节点，不用再写一个循环去找。但是这样也增加了插入删除节点时的维护难度。
四、const_iterator
我们只需要在迭代器的模板参数中多增加两项，即可同时实现普通的迭代器以及 const 版本的迭代器，具体操作如下：
cpp
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;
}
// ...
};那么在使用迭代器的时候，通过传不同的模板参数，即可控制是普通的迭代器还是 const 版本的迭代器。
cpp
// class RBTree
typedef RBTreeIterator<T, T&, T*> Iterator;
typedef RBTreeIterator<T, const T&, const T*> ConstIterator;这时就可以分别在 set 和 map 中定义两个版本的迭代器了。
cpp
// class set
typedef typename RBTree<K, K, SetKeyOfT>::Iterator iterator;
typedef typename RBTree<K, K, SetKeyOfT>::ConstIterator const_iterator;
cpp
// class map
typedef typename RBTree<K, pair<K, V>, MapKeyOfT>::Iterator iterator;
typedef typename RBTree<K, pair<K, V>, MapKeyOfT>::ConstIterator const_iterator;begin() 和 end() 只需加一个 const 即可，其他运算符重载的实现与普通迭代器一样。
cpp
const_iterator begin() const
{
return _t.begin();
}
const_iterator end() const
{
return _t.end();
}五、实现 key 不支持修改
只需在构建红黑树时对应的 key 参数前加上一个 const 即可。
cpp
// class set
RBTree<K, const K, SetKeyOfT> _t;
// class map
RBTree<K, pair<const K, V>, MapKeyOfT> _t;注意其它地方也需要加。
cpp
// class set
typedef typename RBTree<K, const K, SetKeyOfT>::Iterator iterator;
typedef typename RBTree<K, const K, SetKeyOfT>::ConstIterator const_iterator;
// class map
typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::Iterator iterator;
typedef typename RBTree<K, pair<const K, V>, MapKeyOfT>::ConstIterator const_iterator;六、operator[ ]
在 map 中，operator[] 可以实现插入 + 修改的功能。既然有插入，那么就必定要用到 insert 接口。在 STL 库中，insert 的接口的返回值是一个 pair 类型：
- 如果插入成功，返回
<插入的该数据对应的迭代器, true>； - 如果插入失败，返回
<已经存在的该数据对应的迭代器, false>。
之所以这么设计，就是为了实现 operator[]，以下是实现代码：
cpp
V& operator[](const K& key)
{
pair<iterator, bool> ret = insert(make_pair(key, V()));
return ret.first->second;
}那么对应的 insert 函数也应该修改：
cpp
pair<Iterator, bool> Insert(const T& data)
{
if (_root == nullptr)
{
_root = new Node(data);
_root->_col = BLACK;
return { Iterator(_root, _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(cur->_data) > kot(data))
{
parent = cur;
cur = cur->_left;
}
else return { Iterator(cur, _root), false };
}
cur = new Node(data);
Node* newNode = cur;
// ...
return { Iterator(newNode, _root), true };
}七、完整代码
1. RBTree.h
cpp
#pragma once
#include<utility>
#include<iostream>
using namespace std;
// 枚举值表示颜色
enum Colour
{
RED,
BLACK
};
// 由于 set 和 map 底层存储的数据不一样, 可能是 key 也可能是 pair
// 所以红黑树不是写死的, 而是写成一个模板参数 T
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 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;
}
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;
// parent 为空, cur 是根, 说明当前树走完了, 把 nullptr 给 _node, nullptr 作 end()
while (parent && cur == parent->_right)
{
cur = parent;
parent = parent->_parent;
}
_node = parent;
}
return *this;
}
Self& operator--()
{
if (_node == nullptr) // end()
{
// --end()，特殊处理，走到中序最后一个节点，整棵树的最右节点
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;
}
};
// 这里的 K 是 key, 提供给一些需要的函数接口
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 };
}
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 { 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; // 更新 cur 继续向上处理
parent = cur->_parent;
}
// 如果叔叔为黑色或者为空，那么需要旋转 + 变色
else
{
// 如果插入在父亲节点的左边，此时为 LL 型
if (cur == parent->_left)
{
// 右单旋 + 变色
RotateR(grandfather);
parent->_col = BLACK;
grandfather->_col = RED;
}
// LR 型
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
{
// RR 型
if (cur == parent->_right)
{
// 左单旋 + 变色
RotateL(grandfather);
parent->_col = BLACK;
grandfather->_col = RED;
}
// RL 型
else
{
// 右左双旋 + 变色
RotateR(parent);
RotateL(grandfather);
cur->_col = BLACK;
grandfather->_col = RED;
}
// 停止更新
break;
}
}
}
// 变色更新到根节点时有可能根为红色，需要将根变为黑色
_root->_col = BLACK;
return { Iterator(newNode, _root), true };
}
private:
void RotateR(Node* parent)
{
Node* subL = parent->_left; // 左子树
Node* subLR = subL->_right; // 左子树的右子树（b）
parent->_left = subLR;
if (subLR) subLR->_parent = parent; // 注意 b 有可能为空，所以要判断一下
Node* ppnode = parent->_parent;
subL->_right = parent;
parent->_parent = subL;
// parent 有可能是整棵树的根，也可能是局部子树的根
// 如果是整棵树的根，要修改 _root
// 如果是局部的根，指针要跟上一层链接
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;
}
}
private:
Node* _root = nullptr;
};2. myset.h
cpp
#pragma once
#include"RBTree.h"
namespace mine
{
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);
}
private:
RBTree<K, const K, SetKeyOfT> _t;
};
}3. mymap.h
cpp
#pragma once
#include"RBTree.h"
namespace mine
{
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);
}
V& operator[](const K& key)
{
pair<iterator, bool> ret = insert(make_pair(key, V()));
return ret.first->second;
}
private:
RBTree<K, pair<const K, V>, MapKeyOfT> _t;
};
}