题目:普通平衡树

zig zag图示

#include <iostream>
#include <cstring>
#include <algorithm>
using namespace std;
const int N = 100010,INF=0x3f3f3f3f;
struct Node{
    int l,r;//左儿子 右儿子
    int key,val;//key是对应BST的值,是题目给定的数,左<中<右
    //val是对应Heap的值,每个节点的val是随机给的,能保证随机意义下平衡树的高度是O(logn)级别
    int cnt;//对应于求排名有关的量,是这个key的次数
    int size;//是包含自身的子树大小
}tr[N];//平衡数数组,注意的是tr[i]表示i是下标,tr[i]是这个节点
int n,idx,root;//n是操作个数
//idx是每分配一个新的节点时给的下标,实际上在树的存储中都可以这样做
//root是treap的根,之所以要存下来,因为root可能因为后续的左旋和右旋发生改变(这也就是后面涉及到左旋右旋的函数为什么要传引用的原因)
int get_node(int x){//保存一个新的节点
    int p=++idx;
    tr[p].key=x;
    tr[p].val=rand();//treap精髓
    tr[p].size=1;
    tr[p].cnt=1;
    return p;
}
void pushup(int p){
    tr[p].size=tr[tr[p].l].size+tr[tr[p].r].size+tr[p].cnt;
}//类似于线段树的上传,根据子儿子的值更新当前节点的信息
//zig zag 见上图 一定要记得pushup两次(!!!!)
void zig(int& p){
    int q=tr[p].l;
    tr[p].l=tr[q].r,tr[q].r=p,p=q;
    pushup(tr[p].r);
    pushup(p);
}
void zag(int& p){
    int q=tr[p].r;
    tr[p].r=tr[q].l,tr[q].l=p,p=q;
    pushup(tr[p].l);
    pushup(p);
}
//建树函数 注意为了防止多余的讨论,我们建立两个哨兵,一个值是-INF,一个是+INF,但要记得左旋或右旋root(!!!!)
void build(){
    get_node(INF);
    get_node(-INF);
    root=1;
    tr[1].l=2;
    pushup(root);
    if(tr[1].val<tr[2].val)zig(root);
}
//增删的函数其实都遵循着相同的递归范式,是用当前的key和x比较,然后递归到下一层
//插入函数
void insert(int& p,int x){
    if(!p){p=get_node(x);return;}
    if(tr[p].key<x){
        insert(tr[p].r,x);//这是很妙的地方,insert的递归可以理解为是把p的右子树已完成插入和旋转(即满足tr[tr[p].r].val是右子树中最大的)
        if(tr[tr[p].r].val>tr[p].val)zag(p);//再和当前的节点比val,如果它比当前的val大,那就旋转
    }
    else if(tr[p].key==x){
        tr[p].cnt++;
    }
    else{
        insert(tr[p].l,x);
        if(tr[tr[p].l].val>tr[p].val)zig(p);
    }
    pushup(p);//最后记得pushup
}
void remove(int& p,int x){
    if(!p)return ;
    if(tr[p].key>x){
        remove(tr[p].l,x);
    }
    else if(tr[p].key==x){//这里要特别注意,如果当前节点的cnt大于1,那直接减一就好,如果当前节点已经是叶节点,那直接删了就行(对应第二行),不然我们先通过左旋或者右旋把它往下移一层,再递归地删它,最后再pushup一下
        if(tr[p].cnt>1)tr[p].cnt--;
        else if(!tr[p].l&&!tr[p].r)p=0;
        else{
            if(tr[p].l||tr[tr[p].l].val>tr[tr[p].r].val){
                zig(p);
                remove(tr[p].r,x);
            }
            else {
                zag(p);
                remove(tr[p].l,x);
            }
        }
    }
    else{
        remove(tr[p].r,x);
    }
    pushup(p);
}
//下面的四个函数都是静态查找,不需要传引用
int get_rank_by_key(int p,int x){//注意:因为递归的性质,这里的rank是在当前的节点和它的子树里的rank,不含它的祖先节点
    if(!p)return -1;
    if(tr[p].key==x)return tr[tr[p].l].size+1;
    else if(tr[p].key>x)return get_rank_by_key(tr[p].l,x);
    else return get_rank_by_key(tr[p].r,x)+tr[p].cnt+tr[tr[p].l].size;
    
}
int get_key_by_rank(int p,int x){
    if(!p)return INF;
    if(x<=tr[tr[p].l].size)return get_key_by_rank(tr[p].l,x);
    else if(x<=tr[tr[p].l].size+tr[p].cnt)return tr[p].key;
    else return get_key_by_rank(tr[p].r,x-tr[tr[p].l].size-tr[p].cnt);
}
int get_prev(int p,int x){
    if(!p)return -INF;
    if(tr[p].key>=x)return get_prev(tr[p].l,x);
    else return max(tr[p].key,get_prev(tr[p].r,x));
}
int get_next(int p,int x){
    if(!p)return INF;
    if(tr[p].key<=x)return get_next(tr[p].r,x);
    else return min(tr[p].key,get_next(tr[p].l,x));
}
int main(){
    cin>>n;
    build();
    while(n--){
        int op,x;
        cin>>op>>x;
        if(op==1)insert(root,x);
        if(op==2)remove(root,x);
        if(op==3)cout<<get_rank_by_key(root,x)-1<<endl;
        if(op==4)cout<<get_key_by_rank(root,x+1)<<endl;
        if(op==5)cout<<get_prev(root,x)<<endl;
        if(op==6)cout<<get_next(root,x)<<endl;
        
    }
}