二叉平衡树的插入操作 (C++ )
#pragma once
#if 1
using namespace std;
#include<assert.h>
template<class K,class V>
struct AVLTreeNode {
AVLTreeNode<K, V>* _left;
AVLTreeNode<K, V>* _right;
AVLTreeNode<K, V>* _parent;
int _bf;
pair<K, V> _kv;
AVLTreeNode(const pair<K, V>& kv)
: _left(nullptr)
, _right(nullptr)
, _parent(nullptr)
, _kv(kv)
, _bf(0)
{}
};
template<class K,class V>
class AVLTree {
typedef AVLTreeNode<K, V> Node;
public :
bool Insert(const pair<K, V>& kv) {
if (_root == nullptr) {
_root = new Node(kv);
return true;
}
Node* parent = nullptr;
Node* cur = _root;
while (cur) {
if (cur->_kv.first < kv.first) {
parent = cur;
cur = cur->_right;
}else if (cur->_kv.first > kv.first) {
parent = cur;
cur = cur->_left;
}else {
return false;
}
}
cur = new Node(kv);
if (parent->_kv.first < cur->_kv.first) {
parent->_right = cur;
cur->_parent = parent;
}else {
parent->_left = cur;
cur->_parent = parent;
}
while (parent != nullptr) {
if (parent->_left == cur) {
parent->_bf--;
}else {
parent->_bf++;
}
if (parent->_bf == 0) {
break;
}else if (parent->_bf == 1 || parent->_bf == -1) {
cur = parent;
parent = parent->_parent;
}else if (parent->_bf == 2 || parent->_bf == -2) {
if (parent->_bf == 2) {
if (cur->_bf == 1) {
RotateL(parent);
}
else if (cur->_bf == -1) {
RotateRL(parent);
}
}
else if (parent->_bf == -2) {
if (cur->_bf == -1) {
RotateR(parent);
}
else if (cur->_bf == 1) {
RotateLR(parent);
}
}
}
else {
assert(false);
}
}
return true;
}
private :
Node* _root = nullptr;
void RotateL(Node* parent) {
Node* cur = parent->_right;
Node* curL = cur->_left;
parent->_right = curL;
if (curL) {
curL->_parent = parent;
}
cur->_left = parent;
if (parent->_parent) {
if (parent->_parent->_left == parent) {
parent->_parent->_left = cur;
}
else {
parent->_parent->_right = cur;
}
cur->_parent = parent->_parent;
}
else {
_root = cur;
cur->_parent = nullptr;
}
parent->_parent = cur;
parent->_bf = 0;
cur->_bf = 0;
}
void RotateR(Node* parent) {
Node* cur = parent->_left;
Node* curR = cur->_right;
parent->_left = curR;
if (parent->_parent) {
if (parent->_parent->_left == parent) {
parent->_parent->_left = cur;
}
else {
parent->_parent->_right = cur;
}
cur->_parent = parent->_parent;
}
else {
_root = cur;
cur->_parent = nullptr;
}
cur->_right = parent;
if (parent->_parent) {
cur->_parent = parent->_parent;
}
parent->_parent = cur;
parent->_bf = 0;
cur->_bf = 0;
}
void RotateRL(Node* parent) {
Node* cur = parent->_right;
Node* curL = cur->_left;
int bf = curL->_bf;
RotateR(cur);
RotateL(parent);
if (bf == -1) {
parent->_bf = 0;
cur->_bf = 1;
curL->_bf = 0;
}
else if (bf == 1) {
curL->_bf = 0;
cur->_bf = 0;
parent->_bf = -1;
}
else if (bf == 0) {
curL->_bf = 0;
cur->_bf = 0;
parent->_bf = 0;
}
else {
assert(false);
}
}
void RotateLR(Node* parent) {
Node* cur = parent->_left;
Node* curR = cur->_right;
int bf = curR->_bf;
RotateL(cur);
RotateR(parent);
if (bf == -1) {
cur->_bf = -1;
curR->_bf = 0;
parent->_bf = 0;
}
else if (bf == -1) {
parent->_bf = 1;
cur->_bf = 0;
curR->_bf = 0;
}
else if(bf == 0){
parent->_bf = 0;
cur->_bf = 0;
curR->_bf = 0;
}
else {
assert(false);
}
}
};
#endif