352 lines
9.1 KiB
C++
352 lines
9.1 KiB
C++
#pragma once
|
|
#include <iostream>
|
|
|
|
using namespace std;
|
|
|
|
namespace Lenyiin
|
|
{
|
|
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), _bf(0), _kv(kv)
|
|
{
|
|
}
|
|
};
|
|
|
|
template <class K, class V>
|
|
class AVLTree
|
|
{
|
|
private:
|
|
typedef AVLTreeNode<K, V> Node;
|
|
|
|
public:
|
|
bool Insert(const pair<K, V> &kv)
|
|
{
|
|
// 1. 先按搜索树的规则进行插入
|
|
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->_left;
|
|
}
|
|
else if (cur->_kv.first < kv.first)
|
|
{
|
|
parent = cur;
|
|
cur = cur->_right;
|
|
}
|
|
else
|
|
{
|
|
return false;
|
|
}
|
|
}
|
|
|
|
// 找到位置
|
|
cur = new Node(kv);
|
|
if (parent->_kv.first > kv.first)
|
|
{
|
|
parent->_left = cur;
|
|
cur->_parent = parent;
|
|
}
|
|
else
|
|
{
|
|
parent->_right = cur;
|
|
cur->_parent = parent;
|
|
}
|
|
|
|
// 2. 更新平衡因子
|
|
while (parent)
|
|
{
|
|
if (cur == parent->_right)
|
|
{
|
|
parent->_bf++; // 插在右边,平衡因子++
|
|
}
|
|
else
|
|
{
|
|
parent->_bf--; // 插在左边,平衡因子--
|
|
}
|
|
|
|
if (parent->_bf == 0)
|
|
{
|
|
// 说明 parent 所在的子树的高度不变,更新结束
|
|
break;
|
|
}
|
|
else if (parent->_bf == 1 || parent->_bf == -1)
|
|
{
|
|
// 说明 parent 所在的子树的高度变了,继续往上更新
|
|
cur = parent;
|
|
parent = parent->_parent;
|
|
}
|
|
else if (parent->_bf == 2 || parent->_bf == -2)
|
|
{
|
|
// parent 所在的子树出现问题了,需要旋转处理一下
|
|
// 1. 旋转前提是保持它依旧是搜索二叉树
|
|
// 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);
|
|
}
|
|
}
|
|
|
|
// 旋转完成后, parent 所在的树的高度恢复到了,插入节点前高度
|
|
// 如果是子树,对上一层没有影响,更新结束
|
|
break;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
}
|
|
|
|
// 左单旋
|
|
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 (parent == _root)
|
|
{
|
|
_root = subR;
|
|
subR->_parent = nullptr;
|
|
}
|
|
// 2. parent 为根的树只是整棵树中的子树,改变链接关系,那么 subR 要顶替他的位置
|
|
else
|
|
{
|
|
if (ppNode->_left == parent)
|
|
{
|
|
ppNode->_left = subR;
|
|
}
|
|
else
|
|
{
|
|
ppNode->_right = subR;
|
|
}
|
|
subR->_parent = ppNode;
|
|
}
|
|
|
|
parent->_bf = subR->_bf = 0;
|
|
}
|
|
|
|
// 右单旋
|
|
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 (parent == _root)
|
|
{
|
|
_root = subL;
|
|
subL->_parent = nullptr;
|
|
}
|
|
else
|
|
{
|
|
if (ppNode->_left == parent)
|
|
{
|
|
ppNode->_left = subL;
|
|
}
|
|
else
|
|
{
|
|
ppNode->_right = subL;
|
|
}
|
|
subL->_parent = ppNode;
|
|
}
|
|
parent->_bf = subL->_bf = 0;
|
|
}
|
|
|
|
// 右左双旋
|
|
void RotateRL(Node *parent)
|
|
{
|
|
Node *subR = parent->_right;
|
|
Node *subRL = subR->_left;
|
|
int bf = subRL->_bf;
|
|
|
|
RotateR(parent->_right);
|
|
RotateL(parent);
|
|
|
|
if (bf == -1)
|
|
{
|
|
parent->_bf = 0;
|
|
subR->_bf = 1;
|
|
subRL->_bf = 0;
|
|
}
|
|
else if (bf == 1)
|
|
{
|
|
parent->_bf = -1;
|
|
subR->_bf = 0;
|
|
subRL->_bf = 0;
|
|
}
|
|
else if (bf == 0)
|
|
{
|
|
parent->_bf = 0;
|
|
subR->_bf = 0;
|
|
subRL->_bf = 0;
|
|
}
|
|
}
|
|
|
|
// 左右双旋
|
|
void RotateLR(Node *parent)
|
|
{
|
|
Node *subL = parent->_left;
|
|
Node *subLR = subL->_right;
|
|
int bf = subLR->_bf;
|
|
|
|
RotateL(subL);
|
|
RotateR(parent);
|
|
|
|
if (bf == 1)
|
|
{
|
|
parent->_bf = 0;
|
|
subL->_bf = -1;
|
|
subLR->_bf = 0;
|
|
}
|
|
else if (bf == -1)
|
|
{
|
|
parent->_bf = 1;
|
|
subL->_bf = 0;
|
|
subLR->_bf = 0;
|
|
}
|
|
else if (bf == 0)
|
|
{
|
|
parent->_bf = 0;
|
|
subL->_bf = 0;
|
|
subLR->_bf = 0;
|
|
}
|
|
}
|
|
|
|
// 中序遍历
|
|
void _InOrder(Node *root)
|
|
{
|
|
if (root == nullptr)
|
|
{
|
|
return;
|
|
}
|
|
|
|
_InOrder(root->_left);
|
|
cout << root->_kv.first << ":" << root->_kv.second << endl;
|
|
_InOrder(root->_right);
|
|
}
|
|
|
|
void InOrder()
|
|
{
|
|
_InOrder(_root);
|
|
cout << endl;
|
|
}
|
|
|
|
// 获取树的高度
|
|
int Height(Node *root)
|
|
{
|
|
if (root == nullptr)
|
|
{
|
|
return 0;
|
|
}
|
|
|
|
int leftHeight = Height(root->_left);
|
|
int rightHeight = Height(root->_right);
|
|
|
|
return leftHeight > rightHeight ? leftHeight + 1 : rightHeight + 1;
|
|
}
|
|
|
|
// 判断是否是平衡二叉树
|
|
bool _IsBalance(Node *root)
|
|
{
|
|
if (root == nullptr)
|
|
{
|
|
return true;
|
|
}
|
|
|
|
int leftHeight = Height(root->_left);
|
|
int rightHeight = Height(root->_right);
|
|
if (abs(leftHeight - rightHeight) > 1)
|
|
{
|
|
return false;
|
|
}
|
|
|
|
return _IsBalance(root->_left) && _IsBalance(root->_right);
|
|
}
|
|
|
|
bool IsBalance()
|
|
{
|
|
return _IsBalance(_root);
|
|
}
|
|
|
|
// 查找
|
|
Node *Find(const Node *root)
|
|
{
|
|
Node *cur = _root;
|
|
while (cur)
|
|
{
|
|
if (cur->_kv.first > root->_kv.first)
|
|
{
|
|
cur = cur->_left;
|
|
}
|
|
else if (cur->_kv.first < root->_kv.first)
|
|
{
|
|
cur = cur->_right;
|
|
}
|
|
else
|
|
{
|
|
return cur;
|
|
}
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
private:
|
|
Node *_root = nullptr;
|
|
};
|
|
} |