[笔记]AVL树

AVL树是一种严格平衡的二叉搜索树，任何操作结束后，都能保证每个节点的左右子树高度相差不超过\(1\)。
内容源自BV1rt411j7Ff - 【AgOHの数据结构】平衡树专题之叁 树旋转与AVL树。

模板题：P3369 【模板】普通平衡树。

结构体定义 & 基本函数

struct node{
	int l;//左孩子
    int r;//右孩子
    int v;//值
    int hei;//高度，叶节点为1
    int siz;//大小
}avl[N];
int cnt;//当前用到哪一个节点了，用于新建节点
int root;//根节点

//新建节点
void newnode(int &u,int v){
	avl[u=++cnt].v=v;//赋值
	avl[cnt].siz=1;//叶子节点
}

//更新节点信息
void update(int u){
	avl[u].siz=avl[avl[u].l].siz+avl[avl[u].r].siz+1;//左+右+1
	avl[u].hei=max(avl[avl[u].l].hei,avl[avl[u].r].hei)+1;//max(左,右)+1
}

左右旋转

AVL树用旋转来维护树的平衡。旋转分左旋和右旋：

//左旋
void lrot(int &u){
	int r=avl[u].r;
	avl[u].r=avl[r].l;
	avl[r].l=u;
	u=r;
	update(avl[u].l),update(u);
}
//右旋
void rrot(int &u){
	int l=avl[u].l;
	avl[u].l=avl[l].r;
	avl[l].r=u;
	u=l;
	update(avl[u].r),update(u);
}

接下来我们需要判断并处理AVL树的不平衡情况。

//计算平衡因子（即左子树高度-右子树高度）
int factor(int u){
	return avl[avl[u].l].hei-avl[avl[u].r].hei;
}

对于树上的节点\(u\)（假定\(u\)的子树都平衡），其不平衡状态有\(4\)种：

• LL：\(u\)的左子树过高，而左子节点的左子树较高。
  处理方法：右旋一次\(u\)。
• LR：\(u\)的左子树过高，而左子结点的右子树较高。
  处理方法：设\(v\)是\(u\)的左儿子，先左旋\(v\)（转化成LL），再右旋\(u\)。
• RR：\(u\)的右子树过高，而右子节点的右子树较高。
  处理方法：左旋一次\(u\)。
• RL：\(u\)的右子树过高，而右子节点的左子树较高。
  处理方法：设\(v\)是\(u\)的右儿子，先右旋\(v\)（转化成RR），再左旋\(u\)。

若左子节点的左右子树高度相同，则既可以归纳为LL，也可以作为LR考虑。右子节点同理。

//检查并调整为平衡状态，并更新节点的信息
void check(int &u){
	int uf=factor(u);
	if(uf>1){
		int lf=factor(avl[u].l);
		if(lf>0) rrot(u);//LL
		else lrot(avl[u].l),rrot(u);//LR
	}else if(uf<-1){
		int rf=factor(avl[u].r);
		if(rf<0) lrot(u);//RR
		else rrot(avl[u].r),lrot(u);//RL
	}else if(u) update(u);//如果原本就平衡，且u不为空，就要更新
}

其他操作

和普通的BST一样了。

//插入
void ins(int &u,int v){
	if(!u) newnode(u,v);
	else if(v<avl[u].v) ins(avl[u].l,v);
	else ins(avl[u].r,v);
	check(u);//自下向上更新节点信息&调整结构
}

//找u的后继（即u先往右走，再不断往左直到没有左子结点）v，
//让v的父节点直接连接v的右子树
int find(int &u,int fa){
	int ans;
	if(!avl[u].l){//终点
		ans=u;
		avl[fa].l=avl[u].r;
	}else{
		ans=find(avl[u].l,u);
		check(u);
	}
	return ans;
}
//删除
void del(int &u,int v){
	if(v==avl[u].v){
		int l=avl[u].l,r=avl[u].r;
		if(!l||!r) u=l+r;
		else{
			u=find(r,r);//u的后继v来替代u的位置
			avl[u].l=l;//v成为子树的根，连接左边
			if(u!=r) avl[u].r=r;//连接右边
		}
	}else if(v<avl[u].v) del(avl[u].l,v);
	else del(avl[u].r,v);
	check(u);//自下向上更新节点信息&调整结构
}

//计算v的排名（小于v的个数+1）
int getrank(int v){
	int u=root,ran=1;
	while(u){
		if(v<=avl[u].v) u=avl[u].l;
		else{
			ran+=avl[avl[u].l].siz+1;
			u=avl[u].r;
		}
	}
	return ran;
}

//计算第ran名
int getnum(int ran){
	int u=root;
	while(u){
		if(avl[avl[u].l].siz+1==ran) break;
		else if(avl[avl[u].l].siz>=ran)
			u=avl[u].l;
		else
			ran-=avl[avl[u].l].siz+1,u=avl[u].r;
	}
	return avl[u].v;
}

//虽然这种写法可能慢一些，但是它好写
//前驱
int pre(int x){return getnum(getrank(x)-1);}
//后继
int nex(int x){return getnum(getrank(x+1));}

Code

点击查看代码

#include<bits/stdc++.h>
#define int long long
#define N 100010
using namespace std;
struct node{
	int l,r,v,hei,siz;
}avl[N];
int t,cnt,root;
void newnode(int &u,int v){
	avl[u=++cnt].v=v;
	avl[cnt].siz=1;
}
void update(int u){
	avl[u].siz=avl[avl[u].l].siz+avl[avl[u].r].siz+1;
	avl[u].hei=max(avl[avl[u].l].hei,avl[avl[u].r].hei)+1;
}
int factor(int u){
	return avl[avl[u].l].hei-avl[avl[u].r].hei;
}
void lrot(int &u){
	int r=avl[u].r;
	avl[u].r=avl[r].l;
	avl[r].l=u;
	u=r;
	update(avl[u].l),update(u);
}
void rrot(int &u){
	int l=avl[u].l;
	avl[u].l=avl[l].r;
	avl[l].r=u;
	u=l;
	update(avl[u].r),update(u);
}
void check(int &u){
	int uf=factor(u);
	if(uf>1){
		int lf=factor(avl[u].l);
		if(lf>=0) rrot(u);//LL
		else lrot(avl[u].l),rrot(u);//LR
	}else if(uf<-1){
		int rf=factor(avl[u].r);
		if(rf<=0) lrot(u);//RR
		else rrot(avl[u].r),lrot(u);//RL
	}else if(u) update(u);
}
void ins(int &u,int v){
	if(!u) newnode(u,v);
	else if(v<avl[u].v) ins(avl[u].l,v);
	else ins(avl[u].r,v);
	check(u);
}
int find(int &u,int fa){
	int ans;
	if(!avl[u].l){//终点
		ans=u;
		avl[fa].l=avl[u].r;
	}else{
		ans=find(avl[u].l,u);
		check(u);
	}
	return ans;
}
void del(int &u,int v){
	if(v==avl[u].v){
		int l=avl[u].l,r=avl[u].r;
		if(!l||!r) u=l+r;
		else{
			u=find(r,r);//找u的后继，即比u大的第一个数
			avl[u].l=l;
			if(u!=r) avl[u].r=r;
		}
	}else if(v<avl[u].v) del(avl[u].l,v);
	else del(avl[u].r,v);
	check(u);
}
int getrank(int v){
	int u=root,ran=1;
	while(u){
		if(v<=avl[u].v) u=avl[u].l;
		else{
			ran+=avl[avl[u].l].siz+1;
			u=avl[u].r;
		}
	}
	return ran;
}
int getnum(int ran){
	int u=root;
	while(u){
		if(avl[avl[u].l].siz+1==ran) break;
		else if(avl[avl[u].l].siz>=ran)
			u=avl[u].l;
		else
			ran-=avl[avl[u].l].siz+1,u=avl[u].r;
	}
	return avl[u].v;
}
int pre(int x){return getnum(getrank(x)-1);}
int nex(int x){return getnum(getrank(x+1));}
signed main(){
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
	cin>>t;
	while(t--){
		int op,x;
		cin>>op>>x;
		if(op==1) ins(root,x);
		else if(op==2) del(root,x);
		else if(op==3) cout<<getrank(x)<<"\n";
		else if(op==4) cout<<getnum(x)<<"\n";
		else if(op==5) cout<<pre(x)<<"\n";
		else if(op==6) cout<<nex(x)<<"\n";
	}
	return 0;
}

附：相同节点合并写法

当时看完视频，想到是不是能把相同节点计数，存在一个节点中。

于是就写出下面的代码了。结构体多存了一个\(cnt\)，然后newnode、update、ins、del、getrank、getnum函数需要做相应的修改。

点击查看代码

#include<bits/stdc++.h>
#define int long long
#define N 100010
using namespace std;
struct node{
	int l,r,v,hei,siz,cnt;
}avl[N];
int t,cnt,root;
void newnode(int &u,int v){
	avl[u=++cnt].v=v;
	avl[cnt].siz=1;
	avl[cnt].cnt=1;
}
void update(int u){
	avl[u].siz=avl[avl[u].l].siz+avl[avl[u].r].siz+avl[u].cnt;
	avl[u].hei=max(avl[avl[u].l].hei,avl[avl[u].r].hei)+1;
}
int factor(int u){
	return avl[avl[u].l].hei-avl[avl[u].r].hei;
}
void lrot(int &u){
	int r=avl[u].r;
	avl[u].r=avl[r].l;
	avl[r].l=u;
	u=r;
	update(avl[u].l),update(u);
}
void rrot(int &u){
	int l=avl[u].l;
	avl[u].l=avl[l].r;
	avl[l].r=u;
	u=l;
	update(avl[u].r),update(u);
}
void check(int &u){
	int uf=factor(u);
	if(uf>1){
		int lf=factor(avl[u].l);
		if(lf>=0) rrot(u);//LL
		else lrot(avl[u].l),rrot(u);//LR
	}else if(uf<-1){
		int rf=factor(avl[u].r);
		if(rf<=0) lrot(u);//RR
		else rrot(avl[u].r),lrot(u);//RL
	}else if(u) update(u);
}
void ins(int &u,int v){
	if(!u) newnode(u,v);
	else if(v==avl[u].v) avl[u].cnt++;
	else if(v<avl[u].v) ins(avl[u].l,v);
	else ins(avl[u].r,v);
	check(u);
}
int find(int &u,int fa){
	int ans;
	if(!avl[u].l){//终点
		ans=u;
		avl[fa].l=avl[u].r;
	}else{
		ans=find(avl[u].l,u);
		check(u);
	}
	return ans;
}
void del(int &u,int v){
	if(v==avl[u].v){
		if(avl[u].cnt>1) avl[u].cnt--;
		else{
			int l=avl[u].l,r=avl[u].r;
			if(!l||!r) u=l+r;
			else{
				u=find(r,r);//找u的后继，即比u大的第一个数
				avl[u].l=l;
				if(u!=r) avl[u].r=r;
			}
		}
	}else if(v<avl[u].v) del(avl[u].l,v);
	else del(avl[u].r,v);
	check(u);
}
int getrank(int v){//小于自己的个数+1
	int u=root,ran=1;
	while(u){
		if(v<=avl[u].v) u=avl[u].l;
		else{
			ran+=avl[avl[u].l].siz+avl[u].cnt;
			u=avl[u].r;
		}
	}
	return ran;
}
int getnum(int ran){
	int u=root;
	while(u){
		int sz=avl[avl[u].l].siz+avl[u].cnt;
		if(ran<=avl[avl[u].l].siz) u=avl[u].l;
		else if(ran>avl[avl[u].l].siz+avl[u].cnt) ran-=sz,u=avl[u].r;
		else break;//如果ran在[siz[l]+1,siz[l]+cnt[u]]的区间内，就说明第ran名就是u
	}
	return avl[u].v;
}
int pre(int x){return getnum(getrank(x)-1);}
int nex(int x){return getnum(getrank(x+1));}
signed main(){
	ios::sync_with_stdio(false);
	cin.tie(nullptr);
	cin>>t;
	while(t--){
		int op,x;
		cin>>op>>x;
		if(op==1) ins(root,x);
		else if(op==2) del(root,x);
		else if(op==3) cout<<getrank(x)<<"\n";
		else if(op==4) cout<<getnum(x)<<"\n";
		else if(op==5) cout<<pre(x)<<"\n";
		else if(op==6) cout<<nex(x)<<"\n";
	}
	return 0;
}

两种写法效率相当，不合并183ms，合并190ms。
似乎相同节点合并反而更慢？