P3369 【模板】普通平衡树

题面

https://www.luogu.org/problem/P3369

题解

#include<cstdio>
#include<iostream>
#include<ctime>
#include<cstdlib>
using namespace std;
const int INF=1e8;

struct node{
  node *ch[2],*fa;
  int val,pri,cnt,size;
} root;

int n,rank,pre,suc;

bool fsr(node *x){
  return (x->fa->ch[1]!=NULL && x->fa->ch[1]==x);
}

void rotate(node *x){
  int opt;
  node *y=x->fa;
  opt=fsr(x);
  y->ch[opt]=x->ch[1^opt];
  if (x->ch[1^opt]!=NULL) x->ch[1^opt]->fa=y;
  x->fa=y->fa; 
  y->fa->ch[fsr(y)]=x;
  y->fa=x; 
  x->ch[1^opt]=y;

  y->size=y->cnt;
  if (y->ch[0]!=NULL) y->size+=y->ch[0]->size;
  if (y->ch[1]!=NULL) y->size+=y->ch[1]->size;

  x->size=x->cnt;
  x->size+=y->size;
  if (x->ch[opt]!=NULL) x->size+=x->ch[opt]->size;
}

void up(node *x) {
  while (x->pri>x->fa->pri) rotate(x);
}

void down(node* x){
  while (x->ch[0]!=NULL && x->ch[1]!=NULL) {
    if (x->ch[0]->pri>x->ch[1]->pri) rotate(x->ch[0]); 
    else rotate(x->ch[1]);
  }
  if (x->ch[0]==NULL && x->ch[1]==NULL) x->fa->ch[fsr(x)]=NULL;
    else if (x->ch[0]==NULL) x->fa->ch[fsr(x)]=x->ch[1],x->ch[1]->fa=x->fa;
    else x->fa->ch[fsr(x)]=x->ch[0],x->ch[0]->fa=x->fa;
}

void insert(node *now,int x){
  now->size++;
  if (now->val==x) {
    now->cnt++;
    return;
  }
  if (now->ch[now->val>x]==NULL) {
    node *s=new node();
    s->fa=now;
    s->val=x;
    s->pri=rand()%INF;
    s->cnt=1;
    s->size=1;
    now->ch[now->val>x]=s;
    up(s);
    return;
  }
  else insert(now->ch[now->val>x],x);
}

void erase(node *now,int x) {
  if (now==NULL) return;
  now->size--;
  if (now->val==x) {
    now->cnt--;
    if (now->cnt==0) down(now);
  }
  else
    erase(now->ch[x<now->val],x);
}

void dfs(node *x){
  if (x->ch[0]!=NULL) dfs(x->ch[0]);
  printf("val=%d cnt=%d
",x->val,x->cnt);
  if (x->ch[1]!=NULL) dfs(x->ch[1]);
}

void findrank(node *now,int x) {
  if (now->val>x) {
    if (now->ch[1]==NULL) return; else findrank(now->ch[1],x); 
  }
  else {
    if (now->ch[1]!=NULL) rank+=now->ch[1]->size;
    if (now->val==x) {rank++;return;}
    rank+=now->cnt;
    if (now->ch[0]==NULL) return; else findrank(now->ch[0],x);
  }
}

int findbyrank(node *now,int x){
  int le;
  if (now->ch[1]==NULL) le=0; else le=now->ch[1]->size;
  if (rank+le<=x && x<rank+le+now->cnt) return now->val;
  if (x<rank+le) {
    if (now->ch[1]==NULL) return -1; else return findbyrank(now->ch[1],x);
  }
  else {
    rank+=le+now->cnt;
    if (now->ch[0]==NULL) return -1; else return findbyrank(now->ch[0],x);
  }
}

void findpre(node *now,int x){
  if (now->val<x) {
    pre=max(now->val,pre);
    if (now->ch[0]==NULL) return; else findpre(now->ch[0],x);
  }
  else {
    if (now->ch[1]==NULL) return; else findpre(now->ch[1],x);
  }
}

void findsuc(node *now,int x){
  if (now->val>x) {
    suc=min(now->val,suc);
    if (now->ch[1]==NULL) return; else findsuc(now->ch[1],x);
  }
  else {
    if (now->ch[0]==NULL) return; else findsuc(now->ch[0],x);
  }
}

int main(){
  int i,opt,x;
  scanf("%d",&n);
  root.val=-INF; root.pri=INF; root.cnt=1; root.size=1;
  srand(time(0));
  for (i=1;i<=n;i++) {
    scanf("%d %d",&opt,&x);
    if (opt==1) insert(&root,x);
    if (opt==2) erase(&root,x);
    if (opt==3) {
      rank=0;
      findrank(&root,x);
      printf("%d
",rank-1);
    }
    if (opt==4) {
      rank=0;
      printf("%d
",findbyrank(&root,x));
    }
    if (opt==5) {
      pre=-INF;
      findpre(&root,x);
      printf("%d
",pre);
    }
    if (opt==6) {
      suc=INF;
      findsuc(&root,x);
      printf("%d
",suc);
    }
    //dfs(&root);
    //printf("
");
  }
  return 0;
}