COJS：1829. [Tyvj 1728]普通平衡树

  1 #include<bits/stdc++.h>
  2 using namespace std;
  3 const int maxn=100010;
  4 int root,tot,N;
  5 int key[maxn],siz[maxn],lc[maxn],rc[maxn];
  6 inline void r_rotate(int &rt){
  7     int k=lc[rt];
  8     lc[rt]=rc[k];
  9     rc[k]=rt;
 10     siz[k]=siz[rt];
 11     siz[rt]=siz[lc[rt]]+siz[rc[rt]]+1;
 12     rt=k;
 13 }
 14 inline void l_rotate(int &rt){
 15     int k=rc[rt];
 16     rc[rt]=lc[k];
 17     lc[k]=rt;
 18     siz[k]=siz[rt];
 19     siz[rt]=siz[lc[rt]]+siz[rc[rt]]+1;
 20     rt=k;
 21 }
 22 inline void MAINTAIN(int &rt,bool flag){
 23     if(flag==false){//rt的左子树的某子树>rt的右子树 
 24         if(siz[lc[lc[rt]]]>siz[rc[rt]]) r_rotate(rt);//如果rt的左孩子的左孩子>rt的右孩子，rt右旋 
 25         else if(siz[rc[lc[rt]]]>siz[rc[rt]]){//如果rt的左孩子L的右孩子B > rt的右孩子R： 
 26             l_rotate(lc[rt]);//先让rt的左孩子左旋，使rt的左子树变成以B为根 
 27             r_rotate(rt);//再右旋再变成以B为根 
 28         }
 29         else return;//平衡 
 30     }
 31     else{//rt的右子树的某子树>rt的左子树 
 32         if(siz[rc[rc[rt]]]>siz[lc[rt]]) l_rotate(rt);
 33         else if(siz[lc[rc[rt]]]>siz[lc[rt]]){
 34             r_rotate(rc[rt]);
 35             l_rotate(rt);
 36         }
 37         else return;
 38     }
 39     MAINTAIN(lc[rt],0); MAINTAIN(rc[rt],1);//检查是否满足平衡 
 40     MAINTAIN(rt,1); MAINTAIN(rt,0);
 41 }
 42 inline void insert(int &rt,int v){
 43     if(rt==0){//找到合适的位置插入 
 44         rt=++tot; 
 45         key[rt]=v;
 46         lc[rt]=rc[rt]=0; siz[rt]=1;
 47         return ;
 48     }
 49     siz[rt]++;
 50     if(v<=key[rt]) insert(lc[rt],v);
 51     else insert(rc[rt],v);
 52     MAINTAIN(rt,v>key[rt]);//调整树结构 
 53 }
 54 
 55 inline int Delete(int &rt,int v){
 56     int ans;
 57     siz[rt]--;
 58     if(v==key[rt]||(v<key[rt]&&lc[rt]==0)||(v>key[rt]&&rc[rt]==0)){
 59         ans=key[rt];
 60         if(lc[rt]==0||rc[rt]==0) rt=lc[rt]+rc[rt];
 61         else key[rt]=Delete(lc[rt],key[rt]+1);
 62         return ans;
 63     }
 64     if(v<key[rt]) ans=Delete(lc[rt],v);
 65     else ans=Delete(rc[rt],v);
 66     return ans;
 67 }
 68 inline bool find(int &rt,int v){//查找是否有 key值为 v的节点 
 69     if(rt==0) return false;
 70     else if(v==key[rt]) return true;
 71     else if(v<key[rt]) return find(lc[rt],v);
 72     else if(v>key[rt]) return find(rc[rt],v);
 73 }
 74 inline int rank(int &rt,int v){//询问排名 
 75     if(rt==0) return 1;
 76     if(v<=key[rt]) return rank(lc[rt],v);
 77     else return siz[lc[rt]]+1+rank(rc[rt],v);
 78 }
 79 inline int select(int &rt,int k){//返回在v位置上的节点，可以改造成取大和取小函数 
 80     if(k==siz[lc[rt]]+1) return key[rt];
 81     if(k<=siz[lc[rt]]) return select(lc[rt],k);
 82     else return select(rc[rt],k-1-siz[lc[rt]]);
 83 }
 84 inline int pred(int &rt,int v){//返回比 v小的最大的数 
 85     if(rt==0) return v;//要求是比v小，返回v的意思是没找到 
 86     if(v<=key[rt]) return pred(lc[rt],v);//key[rt]>v 必然要往更小的方向，即rt的左子树 
 87     else{//此时 key[rt]<v，而rt的右子树的key值都大于key[rt] 
 88         int ans=pred(rc[rt],v);
 89         if(ans==v) return key[rt];//返回v表示rt的右子树中的key都比v大，返回key[rt] 
 90     }
 91 }
 92 inline int succ(int &rt,int v){//返回比v大的最小的数 
 93     if(rt==0) return v;
 94     if(v>=key[rt]) return succ(rc[rt],v);
 95     else{
 96         int ans=succ(lc[rt],v);
 97         if(ans==v) return key[rt];
 98     }
 99 }
100 inline void inorder(int rt){//输出中序遍历 
101     if(rt!=0){
102         inorder(lc[rt]);
103         cout<<key[rt]<<" ";
104         inorder(rc[rt]);
105     }
106 }
107 int main(){
108     freopen("phs.in","r",stdin);
109     freopen("phs.out","w",stdout);
110     siz[0]=root=tot=0;
111     scanf("%d",&N);
112     for(int i=1,kin,num;i<=N;i++){
113         scanf("%d%d",&kin,&num);
114         if(kin==1) insert(root,num);
115         else if(kin==2) Delete(root,num);
116         else if(kin==3) printf("%d\n",rank(root,num));
117         else if(kin==4) printf("%d\n",select(root,num));
118         else if(kin==5) printf("%d\n",pred(root,num));
119         else if(kin==6) printf("%d\n",succ(root,num));
120     }
121     return 0;
122 }