[模板]洛谷T3369 普通平衡树 链表&普通Treap

Posted 2020-09-23 Running-Coder

  1 #include<cstdio>
  2 #include<iostream>
  3 #include<cstring>
  4 #include<cstdlib>
  5 #include<ctime>
  6 #include<cmath>
  7 #include<string>
  8 #include<stack>
  9 #include<queue>
 10 #include<vector>
 11 #include<map>
 12 #include<algorithm>
 13 using namespace std;
 14 struct node{
 15     int key,yx;  //键值，优先级（小根堆）
 16     int size,num;  //子树元素总数，当前节点元素总数
 17     node *ch[2];  //左右子树指针
 18 };
 19 void rotate(node* &,bool);  //旋转
 20 void insert(node* &,int);  //插入
 21 void del(node* &,int);  //删除
 22 int kth(node *,int);  //求第k大元素
 23 int rank(node *,int);  //求元素排名
 24 void pre(node *,int);  //求前驱
 25 void succ(node *,int);  //求后继
 26 node *root=NULL;
 27 int ans;
 28 int n;
 29 int f,x;
 30 int main(){
 31     srand(time(0));
 32     scanf("%d",&n);
 33     while(n--){
 34         scanf("%d%d",&f,&x);
 35         switch(f){
 36             case 1:insert(root,x);break;
 37             case 2:del(root,x);break;
 38             case 3:printf("%d\n",rank(root,x));break;
 39             case 4:printf("%d\n",kth(root,x));break;
 40             case 5:{
 41                 pre(root,x);
 42                 printf("%d\n",ans);
 43                 break;
 44             }
 45             case 6:{
 46                 succ(root,x);
 47                 printf("%d\n",ans);
 48                 break;
 49             }    
 50         }
 51     }
 52     return 0;
 53 }
 54 void rotate(node* &p,bool f){
 55     node *t=p->ch[f^1];
 56     p->ch[f^1]=t->ch[f];  //改变t的子树的位置
 57     t->ch[f]=p;  //将t旋转至p上方
 58     p->size=p->num;
 59     if(p->ch[0]!=NULL)p->size+=p->ch[0]->size;
 60     if(p->ch[1]!=NULL)p->size+=p->ch[1]->size;
 61     t->size=t->num;
 62     if(t->ch[0]!=NULL)t->size+=t->ch[0]->size;
 63     if(t->ch[1]!=NULL)t->size+=t->ch[1]->size;  //维护节点信息，自底向上先算p再算t，Very important~
 64     p=t;  //将旋转上去的t节点作为当前子树新的根节点，使用引用方式传递
 65 }
 66 void insert(node* &p,int x){
 67     if(p==NULL){
 68         p=(node *)malloc(sizeof(node));
 69         p->key=x;
 70         p->yx=rand();
 71         p->size=p->num=1;
 72         p->ch[0]=p->ch[1]=NULL;
 73         return;
 74     }  //新建节点，在下不知构造函数为何物，只好酱紫了。。。
 75     if(p->key==x){
 76         p->size++;
 77         p->num++;
 78         return;
 79     }  //重复元素直接在相应节点累加元素个数
 80     if(x<p->key){
 81         insert(p->ch[0],x);  //递归插入左子树
 82         if(p->ch[0]->yx<p->yx)rotate(p,1);  //若插入的节点优先级小于当前子树的根节点则将其旋转至根节点上方
 83         else p->size++;  //这样写的正确性证明：1.如果需要旋转，则当前子树的所有节点的信息都已维护完成；2.如果不需旋转，那么在将元素递归插入子树后，只有当前子树的根节点的信息还未更新，那么就将其更新即可
 84     }
 85     else{
 86         insert(p->ch[1],x);
 87         if(p->ch[1]->yx<p->yx)rotate(p,0);
 88         else p->size++;
 89     }  //同理
 90 }
 91 void del(node* &p,int x){
 92     if(p==NULL)return;  //本题的数据保证不会删到空节点，不过写了也无所谓
 93     if(x==p->key){  //当前节点是要删除的元素
 94         if(p->num>1){
 95             p->size--;
 96             p->num--;
 97             return;
 98         }  //若元素重复次数大于1，则减少元素个数即可
 99         else{  //需要删除节点
100             if(p->ch[0]==NULL){
101                 node *t=p;
102                 p=p->ch[1];
103                 free(t);
104                 return;
105             }  //若左子树为空则直接用右子树替代当前节点
106             else if(p->ch[1]==NULL){
107                 node *t=p;
108                 p=p->ch[0];
109                 free(t);
110                 return;
111             }  //若右子树为空则直接用左子树替代当前节点
112             else{  //左右子树均非空，则将当前节点向下旋转，并递归在子树中删除（同时维护小根堆的性质）
113                 if(p->ch[0]->yx<p->ch[1]->yx){  //左子树优先级小于右子树优先级
114                     rotate(p,1);  //将左子树向上旋转，作为当前树新的根节点
115                     del(p->ch[1],x);  //递归在右子树中删除
116                 }
117                 else{
118                     rotate(p,0);
119                     del(p->ch[0],x);
120                 }  //同理
121                 p->size--;  //递归删除完成后，还剩当前根节点的信息未更新，则将其更新即可
122             }
123         }
124     }
125     else{  //当前节点不是要删除的元素
126         if(x<p->key)del(p->ch[0],x);  //递归在左子树中删除
127         else del(p->ch[1],x);  //递归在右子树中删除
128         p->size--;  //递归删除完成后，还剩当前根节点的信息未更新，则将其更新即可
129     }
130 }
131 int kth(node *p,int x){
132     int s=0;  //记录左子树节点数量
133     if(p->ch[0]!=NULL)s=p->ch[0]->size;  //这样写以防RE
134     if(x<=s)return kth(p->ch[0],x);  //查询节点位于左子树内，则其在左子树的排名即是在当前树的排名
135     else if(x<=s+p->num)return p->key;  //当前节点即为所求
136     else return kth(p->ch[1],x-s-p->num);  //查询节点位于右子树内，则其在右子树的排名：当前树排名-左子树元素数-当前树根节点元素数
137 }
138 int rank(node *p,int x){
139     int s=0;
140     if(p->ch[0]!=NULL)s=p->ch[0]->size;  //同理
141     if(x<p->key)return rank(p->ch[0],x);  //查询节点位于左子树内，则其在当前树的排名即是在左子树的排名
142     else if(x==p->key)return s+1;  //找到待查询节点，则其在当前树的排名为：左子树元素数+1
143     else return s+p->num+rank(p->ch[1],x);  //查询节点位于右子树内，则其在当前树的排名为：左子树元素数+当前树根节点元素数+其在右子树的排名
144 }
145 void pre(node *p,int x){
146     if(p==NULL)return;  //递归边界防止RE
147     if(p->key<x){  //当前节点键值为可行解
148         ans=p->key;  //保存可行解
149         pre(p->ch[1],x);  //尝试寻找更优解
150     }
151     else pre(p->ch[0],x);  //当前节点键值不是可行解则回退
152 }
153 void succ(node *p,int x){
154     if(p==NULL)return;
155     if(p->key>x){
156         ans=p->key;
157         succ(p->ch[0],x);
158     }
159     else succ(p->ch[1],x);
160 }  //与pre同理