740 lines
23 KiB
C++
740 lines
23 KiB
C++
#pragma once
|
||
|
||
#include <iostream>
|
||
|
||
namespace Lenyiin
|
||
{
|
||
enum Colour
|
||
{
|
||
RED,
|
||
BLACK
|
||
};
|
||
|
||
template <class T>
|
||
struct RBTreeNode
|
||
{
|
||
RBTreeNode<T> *_left; // 左子节点
|
||
RBTreeNode<T> *_right; // 右子节点
|
||
RBTreeNode<T> *_parent; // 父节点
|
||
|
||
T _data; // 节点存储的数据
|
||
Colour _colour; // 节点的颜色（红色或黑色）
|
||
|
||
// 节点的构造函数，默认为红色节点
|
||
RBTreeNode(const T &data)
|
||
: _left(nullptr), _right(nullptr), _parent(nullptr), _data(data), _colour(RED)
|
||
{
|
||
}
|
||
};
|
||
|
||
template <class T, class Ref, class Ptr>
|
||
struct __TreeIterator
|
||
{
|
||
typedef RBTreeNode<T> Node;
|
||
typedef __TreeIterator<T, Ref, Ptr> Self;
|
||
Node *_node; // 当前迭代器指向的节点
|
||
|
||
// 构造函数
|
||
__TreeIterator(Node *node)
|
||
: _node(node)
|
||
{
|
||
}
|
||
|
||
// 解引用操作符
|
||
Ref operator*()
|
||
{
|
||
return _node->_data;
|
||
}
|
||
|
||
// 访问成员操作符
|
||
Ptr operator->()
|
||
{
|
||
return &_node->_data;
|
||
}
|
||
|
||
// 前置递增运算符
|
||
Self &operator++()
|
||
{
|
||
// 1. 如果右不为空, 中序的下一个就是右子树的最左节点
|
||
if (_node->_right)
|
||
{
|
||
Node *subLeft = _node->_right;
|
||
while (subLeft->_left)
|
||
{
|
||
subLeft = subLeft->_left;
|
||
}
|
||
_node = subLeft;
|
||
}
|
||
// 2. 如果右为空, 表示 _node 所在的子树已经完成, 下一个节点在他祖先中去找, 沿着路径往上找孩子使它的左的那个祖先
|
||
else
|
||
{
|
||
Node *cur = _node;
|
||
Node *parent = cur->_parent;
|
||
while (parent && cur == parent->_right)
|
||
{
|
||
cur = parent;
|
||
parent = cur->_parent;
|
||
}
|
||
_node = parent;
|
||
}
|
||
return *this;
|
||
}
|
||
|
||
// 后置递增运算符
|
||
Self operator++(int)
|
||
{
|
||
Node *tmp = _node;
|
||
++(*this);
|
||
return Self(tmp);
|
||
}
|
||
|
||
// 前置递减运算符
|
||
Self &operator--()
|
||
{
|
||
// 1. 如果左不为空, 中序的上一个就是左树的最右节点
|
||
if (_node->_left)
|
||
{
|
||
Node *subRight = _node->left;
|
||
while (subRight->_right)
|
||
{
|
||
subRight = subRight->_right;
|
||
}
|
||
_node = subRight;
|
||
}
|
||
// 2. 如果左为空, 表示 _node 所在的子树已经完成, 上一个节点在他祖先中去找, 沿着路径往上找孩子是它的右的那个祖先
|
||
else
|
||
{
|
||
Node *cur = _node;
|
||
Node *parent = cur->_parent;
|
||
while (parent && cur == parent->_left)
|
||
{
|
||
cur = parent;
|
||
parent = cur->_parent;
|
||
}
|
||
_node = parent;
|
||
}
|
||
return *this;
|
||
}
|
||
|
||
// 后置递减运算符
|
||
Self operator--(int)
|
||
{
|
||
Node *tmp = _node;
|
||
--(*this);
|
||
return Self(tmp);
|
||
}
|
||
|
||
bool operator!=(const Self &n)
|
||
{
|
||
return _node != n._node;
|
||
}
|
||
|
||
bool operator==(const Self &n)
|
||
{
|
||
return _node == n._node;
|
||
}
|
||
};
|
||
|
||
template <class K, class T, class KOfT>
|
||
class RBTree
|
||
{
|
||
private:
|
||
typedef RBTreeNode<T> Node;
|
||
|
||
// 左单旋
|
||
void RotateL(Node *parent)
|
||
{
|
||
Node *ppNode = parent->_parent;
|
||
Node *subR = parent->_right;
|
||
Node *subRL = subR->_left;
|
||
|
||
parent->_right = subRL;
|
||
if (subRL)
|
||
{
|
||
subRL->_parent = parent;
|
||
}
|
||
subR->_left = parent;
|
||
parent->_parent = subR;
|
||
|
||
// 1. 原来 parent 是这棵树的根, 现在 subR 是根
|
||
if (_root == parent)
|
||
{
|
||
_root = subR;
|
||
subR->_parent = nullptr;
|
||
}
|
||
// 2. parent 为根的树只是整棵树中的子树, 改变链接关系, 那么 subR 要顶替他的位置
|
||
else
|
||
{
|
||
if (ppNode->_left == parent)
|
||
{
|
||
ppNode->_left = subR;
|
||
}
|
||
else
|
||
{
|
||
ppNode->_right = subR;
|
||
}
|
||
subR->_parent = ppNode;
|
||
}
|
||
}
|
||
|
||
// 右单旋
|
||
void RotateR(Node *parent)
|
||
{
|
||
Node *ppNode = parent->_parent;
|
||
Node *subL = parent->_left;
|
||
Node *subLR = subL->_right;
|
||
|
||
parent->_left = subLR;
|
||
if (subLR)
|
||
{
|
||
subLR->_parent = parent;
|
||
}
|
||
subL->_right = parent;
|
||
parent->_parent = subL;
|
||
|
||
if (_root == parent)
|
||
{
|
||
_root = subL;
|
||
subL->_parent = nullptr;
|
||
}
|
||
else
|
||
{
|
||
if (ppNode->_left == parent)
|
||
{
|
||
ppNode->_left = subL;
|
||
}
|
||
else
|
||
{
|
||
ppNode->_right = subL;
|
||
}
|
||
subL->_parent = ppNode;
|
||
}
|
||
}
|
||
|
||
// 删除整个红黑树
|
||
void DeleteTree(Node *root)
|
||
{
|
||
if (root == nullptr)
|
||
{
|
||
return;
|
||
}
|
||
|
||
// 递归删除左子树
|
||
DeleteTree(root->_left);
|
||
// 递归删除右子树
|
||
DeleteTree(root->_right);
|
||
|
||
// 删除当前节点
|
||
delete root;
|
||
}
|
||
|
||
public:
|
||
typedef __TreeIterator<T, T &, T *> iterator;
|
||
typedef __TreeIterator<T, const T &, const T *> const_iterator;
|
||
|
||
iterator begin()
|
||
{
|
||
Node *cur = _root;
|
||
while (cur && cur->_left)
|
||
{
|
||
cur = cur->_left;
|
||
}
|
||
return iterator(cur);
|
||
}
|
||
|
||
iterator end()
|
||
{
|
||
return iterator(nullptr);
|
||
}
|
||
|
||
const_iterator begin() const
|
||
{
|
||
Node *cur = _root;
|
||
while (cur && cur->_left)
|
||
{
|
||
cur = cur->_left;
|
||
}
|
||
return const_iterator(cur);
|
||
}
|
||
|
||
const_iterator end() const
|
||
{
|
||
return const_iterator(nullptr);
|
||
}
|
||
|
||
RBTree(Node *root = nullptr)
|
||
: _root(root)
|
||
{
|
||
}
|
||
|
||
~RBTree()
|
||
{
|
||
DeleteTree(_root);
|
||
}
|
||
|
||
// 插入
|
||
// 1. 空树。插入结点做根，把他变黑。
|
||
// 2. 插入红色节点，他的父亲是黑色的，结束。
|
||
// 3. 插入红色节点，他的父亲是红色的，可以推断他的祖父存在且一定为黑色。关键看叔叔
|
||
// a. 如果叔叔存在且为红，把父亲和叔叔变黑，祖父变红，继续往上处理。
|
||
// b. 如果叔叔存在且为黑，或者不存在。旋转（单旋 or 双旋）+ 变色
|
||
std::pair<iterator, bool> Insert(const T &data)
|
||
{
|
||
// 按照搜索树的规则进行插入
|
||
// 如果树为空，新节点直接作为根节点
|
||
if (_root == nullptr)
|
||
{
|
||
_root = new Node(data);
|
||
_root->_colour = BLACK; // 根节点是黑色的
|
||
return std::make_pair(iterator(_root), true);
|
||
}
|
||
|
||
// 插入时使用比较器来比较键的大小，以支持不同的数据类型
|
||
KOfT koft;
|
||
Node *parent = nullptr;
|
||
Node *cur = _root;
|
||
while (cur)
|
||
{
|
||
if (koft(cur->_data) > koft(data))
|
||
{
|
||
parent = cur;
|
||
cur = cur->_left;
|
||
}
|
||
else if (koft(cur->_data) < koft(data))
|
||
{
|
||
parent = cur;
|
||
cur = cur->_right;
|
||
}
|
||
else
|
||
{
|
||
// 如果键已经存在, 插入无效, 返回 false, 并且返回该键的迭代器
|
||
return std::make_pair(iterator(cur), false);
|
||
}
|
||
}
|
||
|
||
// 找到位置, 根据父节点插入新节点
|
||
cur = new Node(data);
|
||
Node *newnode = cur;
|
||
if (koft(parent->_data) > koft(cur->_data))
|
||
{
|
||
parent->_left = cur;
|
||
cur->_parent = parent;
|
||
}
|
||
else
|
||
{
|
||
parent->_right = cur;
|
||
cur->_parent = parent;
|
||
}
|
||
|
||
// 插入后需要修复红黑树的性质
|
||
InsertFixUp(parent, cur);
|
||
|
||
return std::make_pair(iterator(newnode), true);
|
||
}
|
||
|
||
void InsertFixUp(Node *parent, Node *cur)
|
||
{
|
||
// 调整结点颜色
|
||
// 新结点给红的还是黑的? 红色
|
||
// 1. 空树。插入结点做根, 把他变黑。
|
||
// 2. 插入红色节点, 他的父亲是黑色的, 结束。
|
||
// 3. 插入红色节点, 他的父亲是红色的, 可以推断他的祖父存在且一定为黑色。关键看叔叔
|
||
// a. 如果叔叔存在且为红, 把父亲和叔叔变黑, 祖父变红, 继续往上处理。
|
||
// b. 如果叔叔存在且为黑, 或者不存在。旋转（单旋 or 双旋）+ 变色
|
||
while (parent && parent->_colour == RED)
|
||
{
|
||
// 关键看叔叔
|
||
Node *grandfather = parent->_parent;
|
||
// 父节点是祖父节点的左子节点
|
||
if (grandfather->_left == parent)
|
||
{
|
||
Node *uncle = grandfather->_right;
|
||
|
||
// 情况1: uncle 存在, 且为红
|
||
if (uncle && uncle->_colour == RED)
|
||
{
|
||
parent->_colour = uncle->_colour = BLACK;
|
||
grandfather->_colour = RED;
|
||
|
||
// 继续向上处理
|
||
cur = grandfather;
|
||
parent = cur->_parent;
|
||
}
|
||
// 情况 2 or 情况 3 : uncle 不存在 or uncle 存在且为黑
|
||
else
|
||
{
|
||
// 情况 3 : 双旋 -> 变单旋
|
||
if (cur == parent->_right)
|
||
{
|
||
RotateL(parent);
|
||
std::swap(parent, cur);
|
||
}
|
||
|
||
// 第二种情况 (ps: 有可能是第三种情况变过来的)
|
||
RotateR(grandfather);
|
||
grandfather->_colour = RED;
|
||
parent->_colour = BLACK;
|
||
|
||
break;
|
||
}
|
||
}
|
||
// 父节点是祖父节点的右子节点
|
||
else
|
||
{
|
||
Node *uncle = grandfather->_left;
|
||
|
||
// 情况1: uncle 存在, 且为红
|
||
if (uncle && uncle->_colour == RED)
|
||
{
|
||
parent->_colour = uncle->_colour = BLACK;
|
||
grandfather->_colour = RED;
|
||
|
||
// 继续向上调整
|
||
cur = grandfather;
|
||
parent = cur->_parent;
|
||
}
|
||
// 情况 2 or 情况 3 : uncle 不存在 or uncle 存在且为黑
|
||
else
|
||
{
|
||
// 情况 3 : 双旋 -> 变单旋
|
||
if (cur == parent->_left)
|
||
{
|
||
RotateR(parent);
|
||
std::swap(parent, cur);
|
||
}
|
||
|
||
// 第二种情况 (ps: 有可能是第三种情况变过来的)
|
||
RotateL(grandfather);
|
||
grandfather->_colour = RED;
|
||
parent->_colour = BLACK;
|
||
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
|
||
// 保证根节点为黑
|
||
_root->_colour = BLACK;
|
||
}
|
||
|
||
// 删除操作
|
||
bool Erase(const K &key)
|
||
{
|
||
Node *nodeToDelete = _root;
|
||
KOfT koft;
|
||
|
||
// 1. 寻找要删除的节点
|
||
while (nodeToDelete)
|
||
{
|
||
if (koft(nodeToDelete->_data) > key)
|
||
{
|
||
nodeToDelete = nodeToDelete->_left;
|
||
}
|
||
else if (koft(nodeToDelete->_data) < key)
|
||
{
|
||
nodeToDelete = nodeToDelete->_right;
|
||
}
|
||
else
|
||
{
|
||
break; // 找到了要删除的节点
|
||
}
|
||
}
|
||
|
||
// 节点若不存在, 直接返回 false
|
||
if (nodeToDelete == nullptr)
|
||
{
|
||
return false;
|
||
}
|
||
|
||
// 执行删除操作
|
||
Node *parent, *child;
|
||
// 保存原节点的颜色，以便后续调整
|
||
Colour originalColour = nodeToDelete->_colour;
|
||
|
||
// 2. 处理删除节点的各种情况
|
||
if (nodeToDelete->_left == nullptr)
|
||
{
|
||
// 情况 1：没有左子节点
|
||
child = nodeToDelete->_right;
|
||
parent = nodeToDelete->_parent;
|
||
Transplant(nodeToDelete, nodeToDelete->_right);
|
||
}
|
||
else if (nodeToDelete->_right == nullptr)
|
||
{
|
||
// 情况 2：没有右子节点
|
||
child = nodeToDelete->_left;
|
||
parent = nodeToDelete->_parent;
|
||
Transplant(nodeToDelete, nodeToDelete->_left);
|
||
}
|
||
else
|
||
{
|
||
// 情况 3：有两个子节点
|
||
// 找到右子树中的最小节点（后继节点）
|
||
Node *successor = Minimum(nodeToDelete->_right);
|
||
originalColour = successor->_colour;
|
||
child = successor->_right;
|
||
if (successor->_parent == nodeToDelete)
|
||
{
|
||
parent = successor;
|
||
}
|
||
else
|
||
{
|
||
// 后继节点有右子节点，用它的右子节点替换它
|
||
Transplant(successor, successor->_right);
|
||
successor->_right = nodeToDelete->_right;
|
||
successor->_right->_parent = successor;
|
||
parent = successor->_parent;
|
||
}
|
||
// 用后继节点替换删除节点
|
||
Transplant(nodeToDelete, successor);
|
||
successor->_left = nodeToDelete->_left;
|
||
successor->_left->_parent = successor;
|
||
successor->_colour = nodeToDelete->_colour; // 保持颜色不变
|
||
}
|
||
|
||
// 删除节点
|
||
delete nodeToDelete;
|
||
|
||
// 3. 修复红黑树的性质
|
||
// 如果删除的节点是黑色, 需要进行调整
|
||
if (originalColour == BLACK)
|
||
{
|
||
EraseFixUp(child, parent);
|
||
}
|
||
|
||
return true;
|
||
}
|
||
|
||
Node *Minimum(Node *node)
|
||
{
|
||
while (node->_left != nullptr)
|
||
{
|
||
node = node->_left; // 一直向左走，直到找到最左节点
|
||
}
|
||
return node;
|
||
}
|
||
|
||
void Transplant(Node *u, Node *v)
|
||
{
|
||
// 如果u是根节点，v成为新的根节点
|
||
if (u->_parent == nullptr)
|
||
{
|
||
_root = v;
|
||
}
|
||
// u是左子节点，用v替换它
|
||
else if (u == u->_parent->_left)
|
||
{
|
||
u->_parent->_left = v;
|
||
}
|
||
// u是右子节点，用v替换它
|
||
else
|
||
{
|
||
u->_parent->_right = v;
|
||
}
|
||
|
||
// 连接v与u的父节点
|
||
if (v != nullptr)
|
||
{
|
||
v->_parent = u->_parent;
|
||
}
|
||
}
|
||
|
||
void EraseFixUp(Node *child, Node *parent)
|
||
{
|
||
while (child != _root && (child == nullptr || child->_colour == BLACK))
|
||
{
|
||
if (child == parent->_left)
|
||
{
|
||
Node *brother = parent->_right;
|
||
// 情况1: child 的兄弟节点 brother 是红色的
|
||
if (brother->_colour == RED)
|
||
{
|
||
brother->_colour = BLACK;
|
||
parent->_colour = RED;
|
||
RotateL(parent);
|
||
brother = parent->_right;
|
||
}
|
||
// 情况2: child 的兄弟节点 brother 是黑色的, 且 brother 的两个节点都是黑色的
|
||
if ((brother->_left == nullptr || brother->_left->_colour == BLACK) &&
|
||
(brother->_right == nullptr || brother->_right->_colour == BLACK))
|
||
{
|
||
brother->_colour = RED;
|
||
child = parent;
|
||
parent = child->_parent;
|
||
}
|
||
else
|
||
{
|
||
// 情况3: brother 是黑色的, 并且 brother 的左子节点是红色, 右子节点是黑色
|
||
if (brother->_right == nullptr || brother->_right->_colour == BLACK)
|
||
{
|
||
if (brother->_left)
|
||
{
|
||
brother->_left->_colour = BLACK;
|
||
}
|
||
brother->_colour = RED;
|
||
RotateR(brother);
|
||
brother = parent->_right;
|
||
}
|
||
// 情况4: brother 是黑色的, 并且 brother 的右子节点是红色
|
||
if (brother)
|
||
{
|
||
brother->_colour = parent->_colour;
|
||
parent->_colour = BLACK;
|
||
if (brother->_right)
|
||
{
|
||
brother->_right->_colour = BLACK;
|
||
}
|
||
RotateL(parent);
|
||
}
|
||
child = _root;
|
||
break;
|
||
}
|
||
}
|
||
else
|
||
{
|
||
Node *brother = parent->_left;
|
||
// 情况1: child 的兄弟节点 brother 是红色的
|
||
if (brother->_colour == RED)
|
||
{
|
||
brother->_colour = BLACK;
|
||
parent->_colour = RED;
|
||
RotateR(parent);
|
||
brother = parent->_left;
|
||
}
|
||
// 情况2: child 的兄弟节点 parent 是黑色的, 且 brother 的两个节点都是黑色的
|
||
if ((brother->_left == nullptr || brother->_left->_colour == BLACK) &&
|
||
(brother->_right == nullptr || brother->_right->_colour == BLACK))
|
||
{
|
||
brother->_colour = RED;
|
||
child = parent;
|
||
parent = child->_parent;
|
||
}
|
||
else
|
||
{
|
||
// 情况3: brother 是黑色的, 并且 brother 的右子节点是红色, 左子节点是黑色
|
||
if (brother->_left == nullptr || brother->_left->_colour == BLACK)
|
||
{
|
||
if (brother->_right)
|
||
{
|
||
brother->_right->_colour = BLACK;
|
||
}
|
||
brother->_colour = RED;
|
||
RotateL(brother);
|
||
brother = parent->_left;
|
||
}
|
||
// 情况4: brother 是黑色的, 并且 brother 的左子节点是红色
|
||
if (brother)
|
||
{
|
||
brother->_colour = parent->_colour;
|
||
parent->_colour = BLACK;
|
||
if (brother->_left)
|
||
{
|
||
brother->_left->_colour = BLACK;
|
||
}
|
||
RotateR(parent);
|
||
}
|
||
child = _root;
|
||
break;
|
||
}
|
||
}
|
||
}
|
||
|
||
if (child)
|
||
{
|
||
child->_colour = BLACK;
|
||
}
|
||
}
|
||
|
||
// 查找节点
|
||
iterator &Find(const K &key)
|
||
{
|
||
KOfT koft;
|
||
Node *cur = _root;
|
||
while (cur)
|
||
{
|
||
if (koft(cur->_data) > key)
|
||
{
|
||
cur = cur->_left;
|
||
}
|
||
else if (koft(cur->_data) < key)
|
||
{
|
||
cur = cur->_right;
|
||
}
|
||
else
|
||
{
|
||
return iterator(cur);
|
||
}
|
||
}
|
||
return iterator(nullptr);
|
||
}
|
||
|
||
// 判断是否是红黑树
|
||
bool IsRBTree()
|
||
{
|
||
Node *root = _root;
|
||
|
||
// 空树也是红黑树
|
||
if (root == nullptr)
|
||
{
|
||
return true;
|
||
}
|
||
|
||
// 1. 判断跟是否是黑色的
|
||
if (root->_colour != BLACK)
|
||
{
|
||
std::cout << "根节点不是黑色" << std::endl;
|
||
return false;
|
||
}
|
||
|
||
// 获取任意一条路径上黑色节点的数量
|
||
size_t blackCount = 0;
|
||
Node *cur = _root;
|
||
while (cur)
|
||
{
|
||
if (cur->_colour == BLACK)
|
||
{
|
||
blackCount++;
|
||
}
|
||
cur = cur->_left;
|
||
}
|
||
|
||
// 判断是否满足红黑树的性质, k 用来记录路径中黑色节点的个数
|
||
size_t k = 0;
|
||
return _IsRBTree(_root, k, blackCount);
|
||
}
|
||
|
||
bool _IsRBTree(Node *root, size_t k, size_t blackCount)
|
||
{
|
||
// 走到 nullptr 之后, 判断 k 和 blackCount 是否相等
|
||
if (root == nullptr)
|
||
{
|
||
// 最终黑色节点个数
|
||
if (blackCount != k)
|
||
{
|
||
std::cout << "违反性质四: 每条路径中黑色节点的个数必须相等" << std::endl;
|
||
return false;
|
||
}
|
||
return true;
|
||
}
|
||
|
||
// 统计黑色节点个数
|
||
if (root->_colour == BLACK)
|
||
{
|
||
k++;
|
||
}
|
||
|
||
// 检测当前节点与其父亲节点是否都为红色
|
||
Node *parent = root->_parent;
|
||
if (parent && parent->_colour == RED && root->_colour == RED)
|
||
{
|
||
std::cout << "违反了性质三: 红色节点的孩子必须是黑色" << std::endl;
|
||
return false;
|
||
}
|
||
|
||
return _IsRBTree(root->_left, k, blackCount) && _IsRBTree(root->_right, k, blackCount);
|
||
}
|
||
|
||
private:
|
||
Node *_root;
|
||
};
|
||
} |