改造红黑树
上一章我们已经实现了一个简易的红黑树
本章在其基础上封装 set
和 map
-
将
RBTreeNode
的模板参数改成class T
,表示存放数据的类型 -
将
Rbtree
的模板参数改成class K, class T
,T
表示红黑树的每个结点存放的数据的类型,K
单独表示键值类型。 -
set
存放的是K
类型数据,map
存放的pair<K, V>
类型数据。查找时,set
的K
类型可以直接用来比较,而map
的pair<K, V>
虽然也支持比较,但是比较方式不是我们需要的。
库中 pair
的 <
运算符重载,先比较 first
,first
相同则比较 second
:
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); }
但是我们只需要比较 first
,pair
的比较不符合我们的要求。
- 所以需要仿函数
KeyOfT
,set
和map
分别实现取出自己的K
。 - 另外 insert 返回一个键值对,与库统一
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)
{}
};
// T决定红黑树存什么类型的数据
// set RBTree<K, K>
// map RBTree<K, pair<K, V>>
// KeyOfT 支持取出T对象中key的仿函数
template<class K, class T, class KeyOfT>
class RBTree
{
typedef RBTreeNode<T> Node;
public:
pair<iterator, bool> Insert(const T& data)
{
if (_root == nullptr)
{
_root = new Node(data);
_root->_col = BLACK;
return make_pair(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(cur->_data) > kot(data))
{
parent = cur;
cur = cur->_left;
}
else
{
return make_pair(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 (parent == grandfather->_left) // 父亲在左,叔叔在右
{
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 make_pair(iterator(newnode), true);
}
private:
void RotateL(Node* parent)
{
Node* subR = parent->_right;
Node* subRL = subR->_left;
parent->_right = subRL;
if (subRL) subRL->_parent = parent; // subRL有可能是空树,需要if判断
Node* ppNode = parent->_parent; // 提前记录祖先,后面用来连接新的parent
subR->_left = parent;
parent->_parent = subR;
if (parent == _root) // 如果parent就是根结点,那么新的根是subR
{
_root = subR;
_root->_parent = nullptr;
}
else // 否则需要祖先来连接新的parent(即subR),注意判断左右
{
if (parent == ppNode->_left)
{
ppNode->_left = subR;
}
else
{
ppNode->_right = subR;
}
subR->_parent = ppNode;
}
}
void RotateR(Node* parent)
{
Node* subL = parent->_left;
Node* subLR = subL->_right;
parent->_left = subLR;
if (subLR) subLR->_parent = parent;
Node* ppNode = parent->_parent;
subL->_right = parent;
parent->_parent = subL;
if (parent == _root)
{
_root = subL;
_root->_parent = nullptr;
}
else
{
if (parent == ppNode->_left)
{
ppNode->_left = subL;
}
else
{
ppNode->_right = subL;
}
subL->_parent = ppNode;
}
}
private:
Node* _root;
};
set 实现自己的仿函数,只需要取出key,顺便实现插入
template<class K>
class Set
{
struct SetKeyOfT
{
const K& operator()(const K& key)
{
return key;
}
};
public:
pair<iterator, bool> insert(const K& key)
{
pair<typename RBTree<K, K, SetKeyOfT>::iterator, bool> ret = _t.Insert(key);
return pair<iterator, bool>(iterator(ret.first._node), ret.second);
}
private:
RBTree<K, K, SetKeyOfT> _t;
};
map 的仿函数要取出键值对中的key
template<class K, class V>
class Map
{
struct MapKeyOfT
{
const K& operator()(const pair<K, V>& kv)
{
return kv.first;
}
};
public:
pair<iterator, bool> insert(const pair<K, V>& kv)
{
return _t.Insert(kv);
}
private:
RBTree<K, pair<K, V>, MapKeyOfT> _t;
};
迭代器
下面实现了基本的 *
->
==
!=
运算符重载
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;
}
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 == nullptr)
{
// 找祖先中,孩子是父亲左的父亲
Node* cur = _node;
Node* parent = cur->_parent;
while (parent && parent->_right == cur)
{
cur = cur->_parent;
parent = parent->_parent;
}
_node = parent;
}
else
{
// 右子树的最左结点
Node* subLeft = _node->_right;
while (subLeft->_left)
{
subLeft = subLeft->_left;
}
_node = subLeft;
}
return *this;
}
Self operator++(int)
{
Self tmp(*this);
++(*this);
return tmp;
}
–就是和++反着来,判断当前位置的左子树是否为空
- 非空:左子树的最右结点就是上一个
- 空:向上找,祖先中第一个孩子是右的结点
Self& operator--()
{
if (_node->_left == nullptr)
{
Node* cur = _node;
Node* parent = _node->_parent;
while (parent && parent->_left == cur)
{
cur = cur->_parent;
parent = parent->_parent;
}
_node = parent;
}
else
{
Node* subRight = _node->_left;
while (subRight->_right)
{
subRight = subRight->_right;
}
_node = subRight;
}
return *this;
}
Self operator--(int)
{
Self tmp(*this);
--(*this);
return tmp;
}
接下来在红黑树里面实现 begin
和 end
,begin
是最左结点,end
为空,依然是一个左闭右开。
//RBTree
typedef __RBTreeIterator<T, T&, T*> iterator;
typedef __RBTreeIterator<T, const T&, const T*> const_iterator;
iterator Begin()
{
Node* subLeft = _root;
while (subLeft && subLeft->_left)
{
subLeft = subLeft->_left;
}
return iterator(subLeft);
}
iterator End()
{
return iterator(nullptr);
}
const_iterator Begin() const
{
Node* subLeft = _root;
while (subLeft && subLeft->_left)
{
subLeft = subLeft->_left;
}
return const_iterator(subLeft);
}
const_iterator End() const
{
return const_iterator(nullptr);
}
//Set
typedef typename RBTree<K, K, SetKeyOfT>::const_iterator iterator;
typedef typename RBTree<K, K, SetKeyOfT>::const_iterator const_iterator;
iterator begin() const
{
return _t.Begin();
}
iterator end() const
{
return _t.End();
}
//Map
typedef typename RBTree<K, pair<K, V>, MapKeyOfT>::iterator iterator;
typedef typename RBTree<K, pair<K, V>, MapKeyOfT>::const_iterator const_iterator;
iterator begin()
{
return _t.Begin();
}
iterator end()
{
return _t.End();
}
查找
iterator 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 iterator(cur);
}
}
return End();
}
map::operator[]
V& operator[](const K& key)
{
pair<iterator, bool> ret = insert(make_pair(key, V()));
return ret.first->second;
}
完整代码
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 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 == nullptr)
{
// 找祖先中,孩子是父亲左的父亲
Node* cur = _node;
Node* parent = cur->_parent;
while (parent && parent->_right == cur)
{
cur = cur->_parent;
parent = parent->_parent;
}
_node = parent;
}
else
{
// 右子树的最左结点
Node* subLeft = _node->_right;
while (subLeft->_left)
{
subLeft = subLeft->_left;
}
_node = subLeft;
}
return *this;
}
Self operator++(int)
{
Self tmp(*this);
++(*this);
return tmp;
}
Self& operator--()
{
if (_node->_left == nullptr)
{
Node* cur = _node;
Node* parent = _node->_parent;
while (parent && parent->_left == cur)
{
cur = cur->_parent;
parent = parent->_parent;
}
_node = parent;
}
else
{
Node* subRight = _node->_left;
while (subRight->_right)
{
subRight = subRight->_right;
}
_node = subRight;
}
return *this;
}
Self operator--(int)
{
Self tmp(*this);
--(*this);
return tmp;
}
bool operator!=(const Self& s) const
{
return _node != s._node;
}
bool operator==(const Self& s) const
{
return _node == s._node;
}
};
// T决定红黑树存什么类型的数据
// set RBTree<K, K>
// map RBTree<K, pair<K, V>>
// KeyOfT 支持取出T对象中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*> const_iterator;
iterator Begin()
{
Node* subLeft = _root;
while (subLeft && subLeft->_left)
{
subLeft = subLeft->_left;
}
return iterator(subLeft);
}
iterator End()
{
return iterator(nullptr);
}
const_iterator Begin() const
{
Node* subLeft = _root;
while (subLeft && subLeft->_left)
{
subLeft = subLeft->_left;
}
return const_iterator(subLeft);
}
const_iterator End() const
{
return const_iterator(nullptr);
}
pair<iterator, bool> Insert(const T& data)
{
if (_root == nullptr)
{
_root = new Node(data);
_root->_col = BLACK;
return make_pair(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(cur->_data) > kot(data))
{
parent = cur;
cur = cur->_left;
}
else
{
return make_pair(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 (parent == grandfather->_left) // 父亲在左,叔叔在右
{
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 make_pair(iterator(newnode), true);
}
iterator 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 iterator(cur);
}
}
return End();
}
private:
void RotateL(Node* parent)
{
Node* subR = parent->_right;
Node* subRL = subR->_left;
parent->_right = subRL;
if (subRL) subRL->_parent = parent; // subRL有可能是空树,需要if判断
Node* ppNode = parent->_parent; // 提前记录祖先,后面用来连接新的parent
subR->_left = parent;
parent->_parent = subR;
if (parent == _root) // 如果parent就是根结点,那么新的根是subR
{
_root = subR;
_root->_parent = nullptr;
}
else // 否则需要祖先来连接新的parent(即subR),注意判断左右
{
if (parent == ppNode->_left)
{
ppNode->_left = subR;
}
else
{
ppNode->_right = subR;
}
subR->_parent = ppNode;
}
}
void RotateR(Node* parent)
{
Node* subL = parent->_left;
Node* subLR = subL->_right;
parent->_left = subLR;
if (subLR) subLR->_parent = parent;
Node* ppNode = parent->_parent;
subL->_right = parent;
parent->_parent = subL;
if (parent == _root)
{
_root = subL;
_root->_parent = nullptr;
}
else
{
if (parent == ppNode->_left)
{
ppNode->_left = subL;
}
else
{
ppNode->_right = subL;
}
subL->_parent = ppNode;
}
}
private:
Node* _root;
};
Set.h
#pragma once
#include "RBTree.h"
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;
iterator begin() const
{
return _t.Begin();
}
iterator end() const
{
return _t.End();
}
pair<iterator, bool> insert(const K& key)
{
pair<typename RBTree<K, K, SetKeyOfT>::iterator, bool> ret = _t.Insert(key);
return pair<iterator, bool>(iterator(ret.first._node), ret.second);
}
iterator find(const K& key)
{
return _t.Find(key);
}
private:
RBTree<K, K, SetKeyOfT> _t;
};
Map.h
#pragma once
#include "RBTree.h"
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<K, V>, MapKeyOfT>::iterator iterator;
typedef typename RBTree<K, pair<K, V>, MapKeyOfT>::const_iterator const_iterator;
iterator begin()
{
return _t.Begin();
}
iterator end()
{
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 = insert(make_pair(key, V()));
return ret.first->second;
}
private:
RBTree<K, pair<K, V>, MapKeyOfT> _t;
};