[笔记]普通平衡树(Splay)

大佬的博客

原题链

代码里的注释比较多,还算详细,将就着看吧

#include <iostream>
#include <cstdio>
#include <cmath>
#include <cstring>
#include <algorithm>
#include <limits.h>
using namespace std;
int f[2000010],child[2000010][2],val[2000010],size[2000010];//child[0]left child,child[1]right child
int n,root = 0,tot = 0,cnt[2000010];//cnt是一个节点出现的个数,tot是节点的总数
//size是当前点和它子树的点的总数,val是权值
void update(int x){
    size[x] = size[child[x][0]] + size[child[x][1]] + cnt[x];
}
void rotate (int x){
    int y = f[x],z = f[y],k = (child[y][1] == x);
    child[z][y == child[z][1]] = x;
    f[y] = x;//x和y互换父子关系
    f[x] = z;//z原本是x的祖父节点,现在是父亲节点
	child[y][k] = child[x][k ^ 1]; 
    f[child[x][k ^ 1]] = y;//这个是还没旋转之前x的儿子
	child[x][k ^ 1] = y;//x原来是y的左儿子,现在y就是x的右儿子;x原来是y的右儿子,现在y就是x的左儿子
//旋转过后,x的一个新儿子是y
    update(x);update(y);
}
void splay (int x, int goal){//如果goal==0则代表要转到根节点去
    while(f[x] != goal){
        int y = f[x],z = f[y];
        if(z != goal){
            bool k = (child[z][0] == y) ^ (child[y][0] == x);//如果xyz可以看做一条链(都是父亲的左儿子或右儿子)的话k=0;
            if(k)
				rotate(x); 
			else rotate(y);//一条链上从上往下旋转
        }
        rotate(x);
    }
    if(goal == 0) root = x;//更新根节点
}
void find (int x){
    int u = root;
    if(u == 0) return;//树是空的
    while(child[u][x > val[u]] && x != val[u])//x>val[u]就往右子树走,x == val就代表找到了
        u = child[u][x > val[u]];
    splay(u,0);//找到了就直接把它转到根节点方便后面的操作
}
void insert (int x){
    int u = root,fa = 0;
    while(u != 0 && x != val[u]){
        fa = u; 
		size[u]++;
        u = child[u][x > val[u]];//大于当前位置就找右子树
    }
    if(u != 0) cnt[u] ++;//已经存在过一个一样权值的点了
    else{
        u = ++ tot;//新建节点
        if(fa == 0) root = u;
        else child[fa][x > val[fa]] = u; //称为父亲的儿子
        child[u][1] = child[u][0] = 0; //叶子节点没有儿子
        size[u] = 1;//字数大小包括自己,大小为1
        val[u] = x; f[u] = fa; cnt[u] = 1;
    }
    splay(u,0);//转到根节点,但上面改了子树的大小,所以一定要update
}
int next(int x, int flag){//找前驱(flag=0)就是比当前点小的最大的数;找后继(flag = 1)就是找比自己大的最小的数
    find(x); 
	int u = root;
    if(val[u] > x && flag) return u;//找后继
    if(val[u] < x && !flag) return u;//找前驱
    u = child[u][flag];//如果找前驱就跳到左子树(这样就不会比x大),这样接下来一路跳右子树就行;找后继同理
    while(child[u][flag ^ 1]) //找前驱就在跳到左子树后一直跳右子树,这样才最大;找后继也是一样的
		u = child[u][flag ^ 1];
    return u;//返回位置
}
void del(int x){
    int l = next(x,0),r = next(x,1);
    splay(l, 0); //把前驱转到根节点
    splay(r,l);  //把后继转到根节点的下面
    int u = child[r][0]; //后继的左儿子,转完后后继的左子树只有一个节点就是要删除的节点
	size[r]--;
    if(cnt[u] > 1) {//有多个权值相同的点按题目要求只删一个
        cnt[u]--;
        splay(u,0);
    }
    else child[r][0] = 0;//将这个点清零
}
int check(int x){
    find(x);
    return size[child[root][0]];
}
int sor(int x){//查询第k大
    int u = root;
    if(size[u] < x) //如果整个树的节点数小于要查询的排名就返回不可能
		return -1;
    while(true){//如果排名比当前节点的左子树的节点数量总和还大,就说明目标在右子树里,而且在右子树里的排名是**查询的排名-左子树大小(包括当前节点)**       
        if(x > size[child[u][0]] + cnt[u]){
            x -= size[child[u][0]] + cnt[u];
            u = child[u][1];
        }
        else{
            if(size[child[u][0]] >= x) //目标在左子树里
				u = child[u][0];
            else return u;//如果不在左子树里就只可能是当前的根节点,因为在上面的if中排除了在右子树的可能
        }
    }
}
int main (){
    insert(-INT_MAX);insert(INT_MAX);
    scanf("%d", &n);
    for(int i = 1; i <= n; i ++) {
        int opt, x; scanf("%d%d", &opt, &x);
        if(opt == 1) 
			insert(x);
        if(opt == 2) 
			del(x);
        if(opt == 3) 
			printf("%d\n", check(x));
        if(opt == 4) 
			printf("%d\n", val[sor(x + 1)]);
        if(opt == 5)  
			printf("%d\n", val[next(x, 0)]);
        if(opt == 6)  
			printf("%d\n", val[next(x, 1)]);
    }
    return 0;
}