看了网上三四篇博客,学习了AVL树维护平衡的方式。但感觉他们给出的代码都有一点瑕疵或者遗漏,懂得了思想之后,花了一些时间把他们几篇的长处结合起来,没有使用指针,实现了一下。每个小逻辑功能都抽象成了函数,应该比较好理解,代码逻辑看起来也比较清晰。
下面给出主要的功能插入和删除。至于其他一些没有动到树结构的操作,如查询,求前驱后继等,同其他BST,没有什么特别。
这里顺带一提,下面的代码中,没有维护子树size,如果要求第K小或者名次,可以在upd函数等处添加有关size的维护,之后便可以支持相关查询了。
#include<iostream>
#include<algorithm>
#define de(x) cout<<#x<<" = "<<x<<endl
using namespace std;
const int maxn=1e5+10;
struct AVL
{
int key,h,lc,rc;
}tree[maxn];
int id,root;
inline int newNode(int k,int l,int r)
{
tree[++id].key=k;
tree[id].lc=l;
tree[id].rc=r;
tree[id].h=0;
return id;
}
inline int height(int id)
{
return id ? tree[id].h : 0;
}
inline void upd(int id)
{
if (!id)
return;
int lh=height(tree[id].lc), rh=height(tree[id].rc);
tree[id].h=max(lh, rh)+1;
}
inline int rightRotate(int id)
{
int lc=tree[id].lc;
tree[id].lc=tree[lc].rc;
tree[lc].rc=id;
upd(id);
upd(lc);
return lc;
}
inline int leftRotate(int id)
{
int rc=tree[id].rc;
tree[id].rc=tree[rc].lc;
tree[rc].lc=id;
upd(id);
upd(rc);
return rc;
}
inline int lrRotate(int id)
{
tree[id].lc=leftRotate(tree[id].lc);
return rightRotate(id);
}
inline int rlRotate(int id)
{
tree[id].rc=rightRotate(tree[id].rc);
return leftRotate(id);
}
inline int balance(int id)
{
if (height(tree[id].lc)-height(tree[id].rc) > 1)
{
int lc=tree[id].lc;
if (height(tree[lc].lc) > height(tree[lc].rc))
return rightRotate(id);
else
return lrRotate(id);
}
else if (height(tree[id].rc)-height(tree[id].lc) > 1)
{
int rc=tree[id].rc;
if (height(tree[rc].lc) < height(tree[rc].rc))
return leftRotate(id);
else
return rlRotate(id);
}
return id;
}
int getMax(int id)
{
if (!id)
return 0;
while (tree[id].rc)
id=tree[id].rc;
return id;
}
int getMin(int id)
{
if (!id)
return 0;
while (tree[id].lc)
id=tree[id].lc;
return id;
}
void insert(int& rt, int v)
{
if (!rt)
rt=newNode(v,0,0);
else if (v < tree[rt].key)
insert(tree[rt].lc, v);
else if (v > tree[rt].key)
insert(tree[rt].rc, v);
rt=balance(rt);
upd(rt);
return;
}
void del(int& rt, int v)
{
if (!rt)
return;
if (v < tree[rt].key)
del(tree[rt].lc, v);
else if (v > tree[rt].key)
del(tree[rt].rc, v);
else
{
if (tree[rt].lc&&tree[rt].rc)
{
if (height(tree[rt].lc) > height(tree[rt].rc))
{
int maxId=getMax(tree[rt].lc);
tree[rt].key=tree[maxId].key;
del(tree[rt].lc, tree[maxId].key);
}
else
{
int minId=getMin(tree[rt].rc);
tree[rt].key=tree[minId].key;
del(tree[rt].rc, tree[minId].key);
}
}
else
rt=tree[rt].lc ? tree[rt].lc : tree[rt].rc;
}
rt=balance(rt);
upd(rt);
}