【数据结构】向STL看齐-模拟实现红黑树

目标：回顾模板分离编译，掌握平衡树的左旋和右旋，掌握红黑树插入过程调整逻辑，即旋转和变色。

红黑树的特点：

1、任意一个结点要么是红色，要么是黑色

2、根结点是黑色

3、一条路径上不能出现两个连续的红色结点

4、从根节点到任一叶子节点的黑色结点个数相同

5、NUL结点默认为黑色结点（可忽略）

使用模板分离编译实现红黑树的基本操作：拷贝构造，赋值重载，插入、判空、查找key是否存在、个数、中序打印等。

RBTree.hpp

#ifndef _RBTREE_H_
#define _RBTREE_H_

#include <iostream>
using namespace std;

// 标识红黑树结点颜色
enum Color{RED, BLACK};

// 红黑树结点类型
template<class K, class V>
struct RBTreeNode
{
RBTreeNode<K, V>* _left;   // 左指针域
RBTreeNode<K, V>* _right;  // 右指针域
RBTreeNode<K, V>* _parent; // 指向父亲
K _key;           // 结点的key值
V _val;           // 结点对应的value
Color _color;     // 结点的颜色

// 构造函数
RBTreeNode(const K &key, const V &val, Color color = RED)
: _left(NULL)
, _right(NULL)
, _parent(NULL)
, _key(key)
, _val(val)
, _color(color)
{}
};

// 红黑树
template<class K, class V>
class RBTree
{
public:
typedef RBTreeNode<K, V> Node;
public:
RBTree()
:_root(NULL)
, _size(0)
{}
explicit RBTree(const RBTree & rb);     // 拷贝构造
RBTree& operator= (const RBTree &rb);   // 赋值运算符重载
bool Insert(const K &key, const V &val);// 红黑树的插入
bool IsRBTree();					    // 是否是红黑树
void Inorder();				            // 中序打印
bool Find(const K& key);                // 查找key是否存在
size_t Size();                          // 结点个数
bool IsEmpty();                         // 判断是否为空
~RBTree();                              // 析构函数

private:
Node* _Copy(const Node* pNode, Node* parent); // 拷贝一棵树
void _Inorder(Node *pNode);       // 递归中序遍历
void RotateLeft(Node* parent);    // 左旋
void RotateRight(Node* parent);   // 右旋
void _Clear(Node*& pNode);        // 清空
// 判断是不是红黑树
bool _IsRBTree(Node* pNode, const size_t blackCount, size_t curBlackCount);
bool _Find(Node* pNode, const K& key);
private:
Node* _root;    // 根节点
size_t _size;   // 结点个数
};

#endif /* _RBTREE_H_ */

重点是插入过程的调整

在分析或画图之前要牢记：1.cur为当前结点 2.p为父节点 3.g为祖父结点 4.u为叔叔结点

一、当pCur为根节点时无需旋转调整，直接让根节点的颜色为黑即可，无需进入循环

二、当pCur不为根节点时，调整出现两个红色结点相连的情况，注意此时cur一定是红，p也是红，g为黑，可分为如下情况：

a）. 叔叔结点存在且为红色， 此时将父亲和叔叔结点调整为黑色，祖父变为红色，然后继续循环



b）. 叔叔结点不存在或者为黑色，即 a)的 else情况，主要有如下情形：

1.p是g的左孩子

如果cur是p的左孩子 ， 需要右旋g   --- 简称 左左 右旋
如果cur是p的右孩子 ， 先左旋p， 然后和上述<1>.中第一条情况一样   --- 简称 左右，先左旋后右旋
2. p是g的右孩子
如果cur是p右孩子   ， 需要左旋g  --- 简称 右右 左旋
如果cur是p的左孩子 ， 先右旋p， 然后和上述<2>. 中第一条情况一样   --- 简称 右左，先右旋后左旋

完成后将 g 的颜色改为红色 p的颜色改为黑色

在test.cpp中有测试用例，基本覆盖了上面的所有情况。

RBTree.cpp

#include "RBTree.hpp"

/********************/
/* 以下为 public函数*/
/*******************/

// 拷贝构造
template<class K, class V>
RBTree<K, V>::RBTree(const RBTree & rb)
:_size(rb._size)
{
Node* parent = NULL;
_root = _Copy(rb._root, parent);
}

// 赋值重载
template<class K, class V>
RBTree<K, V>& RBTree<K, V>::operator=(const RBTree &rb)
{
if (*this != &rb)
{
_Clear(_root);
RBTree temp(rb);
_size = temp._size;
std::swap(_root, temp._root);
}
return *this;
}

// 插入元素
template<class K, class V>
bool RBTree<K, V>::Insert(const K &key, const V &val)
{
// 1. 树为空的时候，直接插入
if (NULL == _root)
{
_root = new Node(key, val, BLACK);
return true;
}

// 2. 树不为空
Node* pCur = _root;    // 保存插入的结点
Node* pParent = NULL;  // 保存插入结点的父亲，方便链接

//		1）、 开始找结点位置
while (pCur != NULL)
{
if (pCur->_key < key)
{
pParent = pCur;
pCur = pCur->_right;
}
else if (pCur->_key > key)
{
pParent = pCur;
pCur = pCur->_left;
}
else
return false;
}
//     2）、位置已找到，插入并连接
pCur = new Node(key, val);

if (pParent->_key < key)
{
pParent->_right = pCur;
}
else if (pParent->_key > key)
{
pParent->_left = pCur;
}
pCur->_parent = pParent;

// 3. 判断是否需要调整树，可整合为如下情况
/************************************************************************/
/* 在分析之前要牢记：1.cur为当前结点 2.p为父节点 3.g为祖父结点 4.u为叔叔结点
/* 1. 当pCur为根节点时无需旋转调整，直接让根节点的颜色为黑即可，无需进入循环
/*
/* 2. 当pCur不为根节点时，调整出现两个红色结点相连的情况，注意此时cur一定是红，p也是红，祖父为黑
/*		a）. 叔叔结点存在且为红色， 此时将父亲和叔叔结点调整为黑色，祖父变为红色，然后继续循环
/*		b）. 叔叔结点不存在或者为黑色，即 a)的 else情况
/*         <1>. p是g的左孩子
/*			    如果cur是p的左孩子 ， 需要右旋g
/*              如果cur是p的右孩子 ， 先左旋p， 然后和上述<1>.中第一条情况一样
/*		   <2>. p是g的右孩子
/*				如果cur是p右孩子   ， 需要左旋g
/*              如果cur是p的左孩子 ， 先右旋p， 然后和上述<2>. 中第一条情况一样
/*          完成后将 g 的颜色改为红色 p的颜色改为黑色
/************************************************************************/

while (pCur != _root && pParent->_color == RED)  // 此条件相当于限制了树的高度大于2
{

Node* pGrandfather = pParent->_parent;		 // 保存祖父结点
Node* pUnclue = NULL;					     // 保存叔叔结点

if (pParent == pGrandfather->_left)          // 找到叔叔结点
pUnclue = pGrandfather->_right;
else
pUnclue = pGrandfather->_left;

if (pUnclue != NULL && RED == pUnclue->_color) // 情况2.a
{
pGrandfather->_color = RED;
pParent->_color = BLACK;
pUnclue->_color = BLACK;
}
else
{
if (pParent == pGrandfather->_left) // 左子树
{
if (pCur == pParent->_right)    // 左右
{
RBTree<K, V>::RotateLeft(pParent);
std::swap(pCur, pParent);
}
RBTree<K, V>::RotateRight(pGrandfather);      // 左左
}
else								// 右子树
{
if (pCur == pParent->_left)     // 右左
{
RBTree<K, V>::RotateRight(pParent);
std::swap(pCur, pParent);
}
RBTree<K, V>::RotateLeft(pGrandfather);      // 右右
}
pGrandfather->_color = RED;
pParent->_color = BLACK;
}

pCur = pGrandfather;     // 跳过黑色结点直接开始调整红色结点
pParent = pCur->_parent;
}
_root->_color = BLACK;
_size++;
return true;
}

// 判断是否满足红黑树性质
template<class K, class V>
bool RBTree<K, V>::IsRBTree()
{
if (_root == NULL)
return true;

if (RED == _root->_color)  // 红黑树性质以
return false;

size_t leftBlackCount = 0; // 保存最左路径上黑结点个数
size_t rightBlackCount = 0; // 保存最右路径上黑结点个数
Node* lNext = _root;
Node* rNext = _root;
while (lNext)
{
if (lNext->_color == BLACK)
leftBlackCount++;
lNext = lNext->_left;
}
while (rNext)
{
if (rNext->_color == BLACK)
rightBlackCount++;
rNext = rNext->_right;
}

// 黑结点个数不同
if (leftBlackCount != rightBlackCount)
return false;

size_t curBlackCount = 0; // 保存每走一步黑结点个数

return RBTree<K, V>::_IsRBTree(_root, curBlackCount, curBlackCount);
}

// 中序打印
template<class K, class V>
void RBTree<K, V>::Inorder()
{
cout << "Inorder : ";
RBTree::_Inorder(_root);
cout << endl;
}

// 查找key是否存在
template<class K, class V>
bool RBTree<K, V>::Find(const K& key)
{
return _Find(_root, key);
}

// 结点个数
template<class K, class V>
size_t RBTree<K, V>::Size()
{
return _size;
}

// 判空
template<class K, class V>
bool RBTree<K, V>::IsEmpty()
{
return NULL == _root;
}

// 析构函数
template<class K, class V>
RBTree<K, V>::~RBTree()
{
_Clear(_root);
}

//////////////////////////////////////////////////////////////////////////
/* 以下为类private 成员函数*/
//////////////////////////////////////////////////////////////////////////

// 切记返回值一定要给模板参数。。。。。
// 拷贝
template<class K, class V>
RBTreeNode<K,V>* RBTree<K, V>::_Copy(const Node* rbNode, Node* parent)
{
Node* ptemp = NULL;
if (rbNode != NULL)
{
ptemp = new Node(rbNode->_key, rbNode->_val, rbNode->_color);
ptemp->_parent = parent;
parent = ptemp;
ptemp->_left = _Copy(rbNode->_left, parent);
ptemp->_right = _Copy(rbNode->_right, parent);
}
return ptemp;
}

// 递归中序遍历
template<class K, class V>
void RBTree<K, V>::_Inorder(Node *pNode)
{
if (pNode != NULL)
{
RBTree<K, V>::_Inorder(pNode->_left);
cout << pNode->_val << " ";
RBTree<K, V>::_Inorder(pNode->_right);
}

}

// 左旋
template<class K, class V>
void RBTree<K, V>::RotateLeft(Node* parent)
{
Node* subR = parent->_right;          // 右孩子
Node* pGrangfather = parent->_parent; // 父亲
Node* subRL = subR->_left;			  // 右孩子的左孩子

if (subRL != NULL)
{
subRL->_parent = parent;      // 1.如果subRL存在则改变其父指针域指向parent
}
parent->_right = subRL;           // 2.改变parent的右指针域
parent->_parent = subR;           // 3.改变parent的父指针域
subR->_left = parent;             // 4. 改变subL的左指针域
subR->_parent = pGrangfather;     // 5. 改变subR的父指针域
if (pGrangfather != NULL)
{
if (pGrangfather->_left == parent)
pGrangfather->_left = subR;
else
pGrangfather->_right = subR;
}
else
{
_root = subR;
}

}

// 右旋
template<class K, class V>
void RBTree<K, V>::RotateRight(Node* parent)
{
Node* subL = parent->_left;           // 左孩子
Node* pGrandfather = parent->_parent; // 父亲
Node* subLR = subL->_right;           // 左孩子的右孩子

if (subLR != NULL)
{
subLR->_parent = parent;     // 1. 如果subLR存在改变subLR的父指针域指向parent
}

parent->_left = subLR;			 // 2. 改变parent的左指针域
parent->_parent = subL;			 // 3. 改变parent的父指针域
subL->_right = parent;           // 4. 改变subL的右指针域
subL->_parent = pGrandfather;    // 5. 改变subL的父指针域

if (pGrandfather != NULL)        // 6. 如果parnet存在父节点，则连接否则改变根节点
{
if (pGrandfather->_left == parent)
pGrandfather->_left = subL;
else
pGrandfather->_right = subL;
}
else
{
_root = subL;
}
}

template<class K, class V>
bool RBTree<K, V>::_Find(Node* pNode, const K& key)
{
if (pNode != NULL)
{
if (pNode->_key == key)
return true;
else if (pNode->_key < key)
return _Find(pNode->_left,key);
else
return _Find(pNode->_right, key);
}
return false;
}

// 判断是不是红黑树
template<class K, class V>
bool RBTree<K, V>::_IsRBTree(Node* pNode, const size_t blackCount, size_t curBlackCount)
{
if (NULL == pNode)
return true;

if (pNode->_color == RED && pNode->_parent != NULL && pNode->_parent->_color == RED)
return false;

if (pNode->_color == BLACK)
curBlackCount++;

if (pNode->_left == NULL && pNode->_right == NULL)
if (blackCount == curBlackCount)
return false;

return RBTree<K, V>::_IsRBTree(pNode->_left, blackCount, curBlackCount) &&
RBTree<K, V>::_IsRBTree(pNode->_right, blackCount, curBlackCount);
}

// 释放所有节点
template<class K, class V>
void RBTree<K, V>::_Clear(Node*& pNode)
{
if (pNode != NULL)
{
_Clear(pNode->_left);
_Clear(pNode->_right);
delete pNode;
pNode = NULL;
}
}

test.cpp

#include "RBTree.cpp"
#include "RBTree.hpp"

#define ARRAYSIZE(arr) sizeof(arr)/sizeof(arr[0])

void TestRBTreeInsert()
{
// test1
int array1[] = { 10, 8, 7, 15, 5, 6, 11, 13, 12 };
RBTree<int, int> rb1;
for (size_t idx = 0; idx <ARRAYSIZE(array1); ++idx)
{
rb1.Insert(array1[idx], array1[idx]);
rb1.Inorder();
cout << "插入 " << array1[idx] << " 满足红黑树性质否？ " << boolalpha << rb1.IsRBTree() << endl;
}
cout << "-------------------------------------------------------" << endl;

// test2
int array2[] = { 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16 };
RBTree<int, int> rb2;
for (size_t idx = 0; idx < ARRAYSIZE(array2); ++idx)
{
rb2.Insert(array2[idx], array2[idx]);
rb2.Inorder();
cout << "插入 " << array2[idx] << " 满足红黑树性质否？ " << boolalpha << rb2.IsRBTree() << endl;
}
cout << "-------------------------------------------------------" << endl;

// test3
RBTree<int, int> rb3;
int array3[] = { 15, 14, 13, 12, 11, 10, 9, 8, 7, 6, 5, 1, 3, 2, 1, 0 };

for (size_t idx = 0; idx < ARRAYSIZE(array3); ++idx)
{
rb3.Insert(array3[idx], array3[idx]);
rb3.Inorder();
cout << "插入 " << array3[idx] << " 满足红黑树性质否？ " << boolalpha << rb3.IsRBTree() << endl;
}
}

void TestCopyCon_Assign()
{
// 测试拷贝构造和复制运算符重载
RBTree<string, int> rb;
rb.Insert("张三", 22);
rb.Insert("里斯", 18);
rb.Insert("呵呵", 12);
rb.Inorder();
RBTree<string, int> rbcopy(rb);
rbcopy.Inorder();
}

void TestOtherFun()
{
RBTree<string, int> rb;
cout << rb.Size() << endl;
cout << boolalpha << rb.IsEmpty() << endl;
rb.Insert("张三", 22);
rb.Insert("里斯", 18);
cout << boolalpha << rb.Find("张三") << endl;
cout << boolalpha << rb.Find("不存在") << endl;
rb.Insert("呵呵", 12);
cout << rb.Size() << endl;
rb.Inorder();
cout << boolalpha << rb.IsEmpty() << endl;
rb.~RBTree();
}

int main()
{
//TestRBTreeInsert();
//TestCopyCon_Assign();
TestOtherFun();

return 0;
}