红黑树完整代码
#pragma once
#include <iostream>
using namespace std;
enum Colour
{
BLACK,
RED,
};
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 K, class V>
class RBTree
{
typedef RBTreeNode<pair<K, V>> Node;
public:
RBTree()
:_root(nullptr)
{}
pair<Node*, bool> Insert(const pair<K, V>& kv)
{
if (_root == nullptr)
{
_root = new Node(kv);
_root->_col = BLACK;
return make_pair(_root, true);
}
Node* parent = nullptr;
Node* cur = _root;
while (cur)
{
if (cur->_data.first < kv.first)
{
parent = cur;
cur = cur->_right;
}
else if (cur->_data.first > kv.first)
{
parent = cur;
cur = cur->_left;
}
else
{
return make_pair(cur, false);
}
}
Node* newnode = new Node(kv);
cur = newnode;
cur->_col = RED;
if (parent->_data.first < kv.first)
{
parent->_right = cur;
cur->_parent = parent;
}
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);
grandfather->_col = RED;
parent->_col = BLACK;
}
else
{
RotateLR(grandfather);
cur->_col = BLACK;
grandfather->_col = RED;
}
break;
}
}
else
{
Node* uncle = grandfather->_left;
if (uncle && uncle->_col == RED)
{
uncle->_col = parent->_col = BLACK;
grandfather->_col = RED;
cur = grandfather;
parent = cur->_parent;
}
else
{
if (cur == parent->_right)
{
RotateL(grandfather);
parent->_col = BLACK;
grandfather->_col = RED;
}
else
{
RotateRL(grandfather);
cur->_col = BLACK;
grandfather->_col = RED;
}
break;
}
}
}
_root->_col = BLACK;
return make_pair(cur, true);
}
void RotateL(Node* parent)
{
Node* subR = parent->_right;
Node* subRL = subR->_left;
parent->_right = subRL;
if (subRL)
subRL->_parent = parent;
Node* Pparent = parent->_parent;
parent->_parent = subR;
subR->_left = parent;
if (_root == parent)
{
_root = subR;
subR->_parent = NULL;
}
else
{
if (Pparent)
{
if (Pparent->_left == parent)
{
Pparent->_left = subR;
}
else
{
Pparent->_right = suR;
}
}
subR->_parent = Pparent;
}
}
void RotateR(Node* parent)
{
Node* subL = parent->_left;
Node* subLR = subL->_right;
parent->_left = subLR;
if (subLR)
{
subLR->_parent = parent;
}
Node* Pparent = parent->_parent;
subL->_right = parent;
parent->_parent = subL;
if (parent == _root)
{
_root = subL;
subL->_parent = nullptr;
}
else
{
if (Pparent->_left == parent)
{
Pparent->_left = subL;
}
else
{
Pparent->_right = subL;
}
subL->_parent = Pparent;
}
}
void RotateRL(Node* parent)
{
RotateR(parent->_right);
RotateL(parent);
}
void RotateLR(Node* parent)
{
RotateL(parent->_left);
RotateR(parent);
}
private:
Node* _root;
};