【班子库&数据结构笔记】平衡树入门之AVL树 NOI2004 郁闷的出纳员

题目链接

平衡树是平衡的二叉搜索树，平衡树能够很好地解决暴力BST的一大缺点：某些子树肥大臃肿导致了查询速度的退化。

其中旋转操作是灵魂，能够通过旋转操作做到平衡两个子树的结点数量。当一方比较肥大的时候就会进行zig左旋 / zag右旋 /zigzag双旋/zagzig双旋来平衡；
同时添加size维护子树大小，还有插入的数字大小一样的时候，开一个cnt记录数量。
Code:

#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
#include 
//
using namespace std;
const int INF = 0x3f3f3f3f;//1.06e9大小
const int mod1 = 1e9 +7;
const int mod2 = 998244353;
const int mod3 = 1e9;
const double PI = acos(-1);
const double eps =1e-3;
typedef unsigned long long ull;
typedef long long ll;
typedef unsigned uint;
const int sq5=616991993;
#define ms(x, n) memset(x,n,sizeof(x))
#define debug(x) printf("%d\n",x);
#define pii pair
#define X first
#define Y second
#define pb push_back
#define lsk k<<1
#define rsk k<<1|1
#define lowbit(a) a&(-a)
const int maxn = 1e5+10;
int tt;
struct node
{
     //AVL - tree
	int lc;
	int rc;
	int h;
	int v;//value
	int sz;
	int tot;//cnt域
}tr[maxn];
int rt;
int pos,delta,low,now,leave;
int right_rotate (int k)//zig右旋
{
     
	int t = tr[k].lc;
	tr[k].lc = tr[t].rc;
	tr[t].rc = k;
	tr[k].h = max(tr[tr[k].lc ].h , tr[tr[k].rc ].h)+1;
	tr[t].h = max(tr[tr[t].lc ].h , tr[tr[t].rc ].h)+1;
	tr[k].sz = tr[tr[k].lc ].sz + tr[tr[k].rc ].sz + tr[k].tot;
	tr[t].sz = tr[tr[t].lc ].sz + tr[tr[t].rc ].sz + tr[t].tot;
	return t;
}
int left_rotate(int k)//zag左旋
{
     
	int t = tr[k].rc;
	tr[k].rc = tr[t].lc;
	tr[t].lc = k;
	tr[k].h = max(tr[tr[k].lc ].h , tr[tr[k].rc ].h)+1;
	tr[t].h = max(tr[tr[t].lc ].h , tr[tr[t].rc ].h)+1;
	tr[k].sz = tr[tr[k].lc ].sz + tr[tr[k].rc ].sz + tr[k].tot;
	tr[t].sz = tr[tr[t].lc ].sz + tr[tr[t].rc ].sz + tr[t].tot;
	return t;
}
int right_left_rotate(int k)//zigzag双旋
{
     
	tr[k].rc = right_rotate(tr[k].rc);
	return left_rotate(k);
}
int left_right_rotate(int k)//zagzig双旋
{
     
	tr[k].lc = left_rotate(tr[k].lc);
	return right_rotate(k);
}
void maintain(int &k)
{
     
	if(tr[tr[k].lc ].h == tr[tr[k].rc ].h+2)
	{
     //左子树高了
		int t=tr[k].lc;
		if(tr[tr[t].lc ].h == tr[tr[k].rc ].h+1)
			k = right_rotate(k);//左子树的左儿子
		else if(tr[tr[t].rc ].h == tr[tr[k].rc ].h+1)
			k = left_right_rotate(k);
	}
	else if(tr[tr[k].rc ].h == tr[tr[k].lc ].h+2)
	{
     
		int t = tr[k].rc;
		if(tr[tr[t].rc ].h == tr[tr[k].lc ].h+1)
			k = left_rotate(k);
		else if(tr[tr[t].lc ].h == tr[tr[k].lc ].h+1)
			k = right_left_rotate(k);
	}
	tr[k].h = max(tr[tr[k].lc ].h , tr[tr[k].rc ].h)+1;
	tr[k].sz = tr[tr[k].lc ].sz + tr[tr[k].rc ].sz + tr[k].tot;
}
int _insert(int k,int x)
{
     
	if(k == 0)
	{
     
		tr[++pos].h = 1;
		tr[pos].v = x;
		tr[pos].sz = 1;
		tr[pos].tot = 1;
		return pos;
	}
	if(x < tr[k].v)tr[k].lc = _insert(tr[k].lc,x);
	else if(x > tr[k].v)tr[k].rc = _insert(tr[k].rc,x);
	else tr[k].sz++ ,tr[k].tot++;
	maintain(k);
	return k;
}
int _rank(int k,int x)
{
     
	if(x>=tr[tr[k].rc ].sz+1&&x<=tr[tr[k].rc ].sz+tr[k].tot)return tr[k].v;
	if(x<tr[tr[k].rc ].sz+1)return _rank(tr[k].rc,x);
	else return _rank(tr[k].lc , x-tr[tr[k].rc ].sz-tr[k].tot);
}
int least_money(int k)
{
     //返回当前最少的工资
	int res = INF;
	while(k)
	{
     
		res = tr[k].v;
		k = tr[k].lc;
	}
	return res;
}
pii delt(int &k,int x)
{
     
	int tv ,ttot;
	pii t;
	if(x == tr[k].v||(x < tr[k].v&&tr[k].lc == 0)||(x > tr[k].v&&tr[k].rc == 0))
	{
     
		if(tr[k].lc == 0||tr[k].rc == 0)
		{
     
			tv =tr[k].v;
			ttot = tr[k].tot;
			k = tr[k].lc + tr[k].rc;
			return make_pair(tv,ttot);
		}
		else
		{
     
			t = delt(tr[k].lc , x);//左子树的最大值代替
			tr[k].v = t.X;
			tr[k].tot = t.Y;
		}
	}
	else
	{
     
		if(x < tr[k].v)t = delt(tr[k].lc,x);
		else t = delt(tr[k].rc,x);
	}
	maintain(k);
	return t;
}
int main()
{
     
	int n,low;
	scanf("%d%d",&n,&low);
	while(n--)
	{
     
		char op;
		int x;
		scanf(" %c %d",&op,&x);
		if(op=='I'&&x>=low)
		{
     
			rt = _insert(rt,x-delta);
			now++;
		}
		else if(op=='A')delta += x;
		else if(op=='S')
		{
     
			int tx;
			pii t;
			delta -= x;
			while( (tx = least_money(rt) )+delta < low)
			{
     
				t = delt(rt,tx);
				now -= t.Y;
				leave += t.Y;
			}
		}
		else if(op=='F')
		{
     
			if(now>=x)printf("%d\n",_rank(rt,x)+delta);
			else printf("-1\n");
		}
	}
	printf("%d\n",leave);
	return 0;
}