【模板】平衡树：FHQ-Treap

posted on 2022-07-21 14:10:58 | under 模板 | source

typedef long long LL;
template<int N> struct fhqtreap{
	int val[N+10],pri[N+10],siz[N+10],x,y,z,root;//小根堆 
	int ch[N+10][2],tot,tsh[N+10],cnt;
	fhqtreap():tot(0),cnt(0){srand(time(0)),siz[0]=0;}
	int newnode(int v=0){int p=cnt?tsh[cnt--]:++tot;return val[p]=v,pri[p]=rand(),siz[p]=1,ch[p][0]=ch[p][1]=0,p;} 
	void destroy(int p){tsh[++cnt]=p;}
	int maintain(int p){return siz[p]=siz[ch[p][0]]+1+siz[ch[p][1]],p;}
	int merge(int p,int q){
		if(!p||!q) return p+q;
		if(pri[p]<pri[q]) return ch[p][1]=merge(ch[p][1],q),maintain(p);
		else return ch[q][0]=merge(p,ch[q][0]),maintain(q);
	}
	void split(int p,int v,int &x,int &y){
		if(!p) return void(x=y=0);
		if(val[p]<=v) x=p,split(ch[p][1],v,ch[p][1],y);//左子树和根在 x
		else split(ch[p][0],v,x,ch[p][0]),y=p;//右子树和根在 y 
		maintain(p); 
	}
	void insert(int v){split(root,v,x,y),root=merge(merge(x,newnode(v)),y);}
	void erase(int v){split(root,v-1,x,y),split(y,v,y,z),destroy(y),root=merge(merge(x,ch[y][0]),merge(ch[y][1],z));}
	int getrnk(int v){split(root,v-1,x,y);int res=siz[x]+1;return root=merge(x,y),res;}
	int getkth(int k,int p){int d=k-siz[ch[p][0]]-1;return !d?val[p]:(d>0?getkth(d,ch[p][1]):getkth(k,ch[p][0]));}
	int getpre(int v){return getkth(getrnk(v)-1,root);}
	int getsuf(int v){return getkth(getrnk(v+1),root);}
};

由这份代码得，fhqtreap 是一种每个函数不超过三行的数据结构，快来写 fhqtreap。

#include <random>
#include <cstdio>
#include <vector>
#include <cstring>
#include <algorithm>
using namespace std;
#ifdef LOCAL
#define debug(...) fprintf(stderr,##__VA_ARGS__)
#else
#define debug(...) void(0)
#endif
typedef long long LL;
mt19937 rng(random_device{}());
template<int N> struct fhqtreap{
	int ch[N+10][2],tot,val[N+10],pri[N+10],siz[N+10],x,y,z,root;
	int newnode(int v){int p=++tot;return ch[p][0]=ch[p][1]=0,pri[p]=rng(),val[p]=v,siz[p]=1,p;}
	void maintain(int p){siz[p]=siz[ch[p][0]]+1+siz[ch[p][1]];}
	fhqtreap():tot(-1),root(0){newnode(0),siz[0]=0;}
	int merge(int p,int q){
		if(!p||!q) return p+q;
		if(pri[p]<pri[q]) return ch[p][1]=merge(ch[p][1],q),maintain(p),p;
		else return ch[q][0]=merge(p,ch[q][0]),maintain(q),q;
	}
	void split(int p,int v,int &x,int &y){
		if(!p) return x=y=0,void();
		if(val[p]<=v) x=p,split(ch[p][1],v,ch[p][1],y),maintain(p);
		else split(ch[p][0],v,x,ch[p][0]),y=p,maintain(p);
	}
	int find(int v,int p){
		if(val[p]==v) return p;
		return val[p]<=v?find(v,ch[p][1]):find(v,ch[p][0]);
	}
	void insert(int v){
		split(root,v,x,y);
		root=merge(x,merge(newnode(v),y));
	}
	void erase(int v){
		split(root,v-1,x,y),split(y,v,y,z);
		root=merge(merge(x,ch[y][0]),merge(ch[y][1],z));
	}
	int getrnk(int v){
		split(root,v-1,x,y); int res=siz[x]+1;
		root=merge(x,y); return res;
	}
	int getkth(int k,int p){
		int c=siz[ch[p][0]]+1;
		if(c==k) return val[p];
		else return c<k?getkth(k-c,ch[p][1]):getkth(k,ch[p][0]);
	}
	int getpre(int v){
		split(root,v-1,x,y);
		int p=x; while(ch[p][1]) p=ch[p][1];
		root=merge(x,y); return val[p];
	}
	int getsuf(int v){
		split(root,v,x,y);
		int p=y; while(ch[p][0]) p=ch[p][0];
		root=merge(x,y); return val[p];
	}
};
fhqtreap<1<<17> t;
int main(){
//	#ifdef LOCAL
//	 	freopen("input.in","r",stdin);
//	#endif
	scanf("%*d");
	for(int op,x;~scanf("%d%d",&op,&x);){
		switch(op){
			case 1: t.insert(x); break;
			case 2: t.erase(x); break;
			case 3: printf("%d\n",t.getrnk(x)); break;
			case 4: printf("%d\n",t.getkth(x,t.root)); break;
			case 5: printf("%d\n",t.getpre(x)); break;
			case 6: printf("%d\n",t.getsuf(x)); break;
		}
	}
	return 0;
}