【Treap模板详细注释】BZOJ3224-普通平衡树

  1 #include<iostream>
  2 #include<cstdio>
  3 #include<cstring>
  4 #include<algorithm>
  5 #include<cmath>
  6 using namespace std;
  7 const int INF=0x7fffffff;
  8 struct treap
  9 {
 10     treap* lson;
 11     treap* rson;
 12     int key;//该节点的值 
 13     int priority;//优先级
 14     int size;//该节点以及其子树的大小
 15     int cnt;//当前值的数字的个数 
 16     int lsize()//返回左子树的大小 
 17     {
 18         if (lson==NULL) return 0;
 19             else return lson->size;
 20     }
 21     int rsize()//返回右子树的大小 
 22     {
 23         if (rson==NULL) return 0;
 24             else return rson->size;
 25     } 
 26     
 27     treap()
 28     {
 29         size=cnt=1;//这里初始值应该是1而不是0 
 30         priority =rand();//随机生成优先级
 31         lson=rson=NULL; 
 32     }
 33 }; 
 34 treap* root=NULL;
 35 int n;
 36 
 37 void update(treap* &rt)//【【Attention】】要写成*& 
 38 {
 39 /*旋转后进行的操作更新size大小=cnt+左子树size+右子树size*/
 40     rt->size=rt->cnt;
 41     rt->size+=rt->lsize();
 42     rt->size+=rt->rsize(); 
 43 } 
 44 
 45 void RightRotate(treap* &rt)
 46 {
 47 /*画图模拟一下:D*/ 
 48     treap* tmp=rt->lson;
 49     rt->lson=tmp->rson;
 50     tmp->rson=rt;
 51     update(rt);
 52     update(tmp);
 53     rt=tmp;
 54 }
 55 
 56 void LeftRotate(treap* &rt)
 57 {
 58     treap* tmp=rt->rson;
 59     rt->rson=tmp->lson;
 60     tmp->lson=rt;
 61     update(rt);
 62     update(tmp);
 63     rt=tmp;
 64 }
 65 
 66 void insert(treap*&rt,int x)
 67 {
 68 /*
 69 如果当前根节点是空的，那么新建一棵平衡树；
 70 如果当前的x和根节点的数值大小一样，则直接在根节点累加个数；
 71 如果当前x小于根节点的数值，那么向左子树搜索。回溯后如果左边的优先级大于根节点，则右旋
 72 如果当前x小于根节点的数值，那么向右子树搜索。回溯后如果右边的优先级大于根节点，则左旋
 73 */ 
 74     if (rt==NULL)
 75     {
 76         rt=new treap;
 77         rt->key=x;
 78     } 
 79         else if (x==rt->key)
 80         {
 81             rt->size++;
 82             rt->cnt++;
 83         }
 84             else if (x<rt->key)
 85             {
 86                 insert(rt->lson,x);
 87                 if (rt->lson->priority>rt->priority) RightRotate(rt);
 88                 update(rt);
 89             }
 90             else
 91             {
 92                 insert(rt->rson,x);
 93                 if (rt->rson->priority>rt->priority) LeftRotate(rt);
 94                 update(rt);
 95             }
 96 } 
 97 
 98 void del(treap* &rt,int x)
 99 {
100 /*
101 如果当前根节点的key就是要删除的那个值
102 （1）如果当前数值的个数大于1个，直接减去cnt和size即可； 
103 （2）如果左子树是空的，这将根节点保存给临时变量，根节点等于它的右子树，释放原来根节点内存； 
104 （3）右子树为空同理； 
105 （4）如果左右子树均不为空，比较两者的优先级。如果左子树优先级大则右旋，反之左旋。
106 如果当前根节点的key不是要删除的那个值
107 （1）如果x<key,则搜索左子树；
108 （2） 否则搜索右子树。 
109 */ 
110     if (x==rt->key)
111     {
112         if (rt->cnt>1)
113         {
114             rt->cnt--;
115             rt->size--;
116         }
117         else if (rt->lson==NULL)
118         {
119             treap* tmp=rt;
120             rt=rt->rson;
121             delete tmp;
122         }
123         else if (rt->rson==NULL)
124         {
125             treap* tmp=rt;
126             rt=rt->lson;
127             delete tmp;
128         }
129         else
130         {
131             if (rt->lson->priority>rt->rson->priority)
132             {
133                 RightRotate(rt);
134                 del(rt->rson,x);
135             }
136             else
137             {
138                 LeftRotate(rt);
139                 del(rt->lson,x);
140             }
141         }
142     }
143     else
144     {
145         if (rt->key>x) del(rt->lson,x);
146             else del(rt->rson,x);
147     }
148     if (rt!=NULL) update(rt);
149 }
150 
151 int getrank(treap* &rt,int x)
152 {
153 /*
154 如果当前的key等于要询问的数，直接返回排名（左子树的大小+1）
155 如果当前的key大于要询问的数，进入左子树查询并返回排名 
156 如果当前的key小于要询问的树，则返回的排名=（左子树大小+当前节点数的个数+它在右子树中的排名） 
157 */
158     if (rt->key==x) return (rt->lsize()+1);
159     if (x<rt->key) return (getrank(rt->lson,x));
160     return (getrank(rt->rson,x)+rt->lsize()+rt->cnt);
161 }
162 
163 int getnum(treap* &rt,int x)   
164 {
165 /*
166 如果需要询问的排名属于当前节点的范围(左子树个数+1~左子树个数+当前节点数的个数)，则直接返回当前的key
167 如果要询问的排名在左子树，则前往左子树询问排名
168 如果要询问的排名在右子树，则前往右子树询问排名，且要询问的排名=排名-左子树的大小-当前节点代表的数的个数 
169 */ 
170     if (rt->lsize()+1<=x && x<=rt->lsize()+rt->cnt) return rt->key;
171     if (x<=rt->lsize()) return getnum(rt->lson,x);
172     return getnum(rt->rson,x-rt->lsize()-rt->cnt);
173 }
174 
175 int pre(treap* &rt,int x)
176 /*
177 如果当前节点的值小于x，则取它和ans中较大的那个，并继续搜索右子树
178 如果当前节点的值大于等于x，则继续搜索左子树 
179 */ 
180 {
181     int ans=-INF;
182     treap* tmp=rt;
183     while (tmp)
184     {
185         if (tmp->key<x)
186         /*尴尬，这整个子程序的tmp->key都手残敲成了tmp->key...下次注意:D*/ 
187         {
188             ans=max(tmp->key,ans);
189             tmp=tmp->rson;
190         } 
191         else tmp=tmp->lson;
192     }
193     return ans;
194 }
195 
196 int suc(treap* &rt,int x)
197 /*与pre类似*/
198 {
199     int ans=INF;
200     treap* tmp=rt;
201     while (tmp)
202     {
203         if (tmp->key>x)
204         {
205             ans=min(tmp->key,ans);
206             tmp=tmp->lson;
207         }
208         else tmp=tmp->rson;
209     }
210     return ans;
211 }
212 
213 void option()
214 {
215     int opt,x;
216     scanf("%d%d",&opt,&x);
217     if (opt==1) insert(root,x);
218     else if (opt==2) del(root,x);
219     else if (opt==3) printf("%d\n",getrank(root,x));
220     else if (opt==4) printf("%d\n",getnum(root,x));
221     else if (opt==5) printf("%d\n",pre(root,x));
222     else if (opt==6) printf("%d\n",suc(root,x));  
223 }
224 
225 void release(treap* &rt)
226 {
227     if (rt->lson) release(rt->lson);
228     if (rt->rson) release(rt->rson);
229     free(rt);
230 } 
231 
232 int main()
233 {
234     scanf("%d",&n);
235     for (int i=0;i<n;i++) option(); 
236     release(root);
237     return 0;
238 }