Treap实现山寨set

  1 #include <iostream>
  2 #include <cstring>
  3 #include <cstdlib>
  4 #include <ctime>
  5 #include <cstdio>
  6 using namespace std;
  7 
  8 struct Node
  9 {
 10     Node* ch[2];
 11     int r,v,s;    //r:随机优先级，v:值，s:以本节点为根的子树的结点总数
 12     bool operator < (const Node& rhs)
 13     {
 14         return (r<rhs.r);
 15     }
 16     int cmp(int x)
 17     {
 18         if (x==v) return -1;
 19         return x<v?0:1;
 20     }
 21     void maintain()
 22     {
 23         s=1;
 24         if (ch[0]!=NULL) s+=ch[0]->s;
 25         if (ch[1]!=NULL) s+=ch[1]->s;
 26     }
 27     Node (int v):v(v)        //结构体的构造函数
 28     {
 29         ch[0]=ch[1]=NULL;
 30         r=rand();
 31         s=1;
 32     }
 33 };
 34 
 35 void rotate(Node* &o,int d)
 36 {
 37     Node* k=o->ch[d^1];
 38     o->ch[d^1]=k->ch[d];
 39     k->ch[d]=o;
 40     o->maintain();
 41     k->maintain();
 42     o=k;
 43 }
 44 
 45 void insert(Node* &o,int x)
 46 {
 47     if (o==NULL)
 48         o=new Node(x);
 49     else
 50     {
 51         //int d=o->cmp(x);        //不用cmp
 52         int d=(x < o->v ? 0 : 1);
 53         insert(o->ch[d],x);
 54         if (o->ch[d]->r > o->r)
 55             rotate(o,d^1);
 56     }
 57     o->maintain();
 58 }
 59 
 60 void remove(Node* &o,int x)
 61 {
 62     int d=o->cmp(x);
 63     if (d==-1)
 64     {
 65         Node *u=o;
 66         if ((o->ch[0]!=NULL) && (o->ch[1]!=NULL))
 67         {
 68             int d2=(o->ch[0]->r > o->ch[1]->r ? 1 : 0);
 69             rotate(o,d2);
 70             remove(o->ch[d2],x);
 71         }
 72         else
 73         {
 74             if (o->ch[0]==NULL)
 75                 o=o->ch[1];
 76             else
 77                 o=o->ch[0];
 78             delete u;
 79         }
 80     }
 81     else
 82         remove(o->ch[d],x);
 83     if (o!=NULL)
 84         o->maintain();
 85 }
 86 
 87 int kth(Node* o,int k)
 88 {
 89     if ((o==NULL)||(k<=0)||(k > o->s))
 90         return 0;
 91     int s=(o->ch[1]==NULL ? 0 : o->ch[1]->s);
 92     if (k==s+1)
 93         return o->v;
 94     else if (k<=s)
 95         return kth(o->ch[1],k);
 96     else
 97         return kth(o->ch[0],k-s-1);
 98 
 99 }
100 
101 int _find(Node* o,int x)
102 {
103     while (o!=NULL)
104     {
105         int d=o->cmp(x);
106         if (d==-1)
107             return 1;
108         else
109             o=o->ch[d];
110     }
111     return 0;
112 }
113 
114 int main()
115 {
116     //freopen("in.txt","r",stdin);
117 
118     int n,m,opr,x;
119     srand(time(0));
120     cin>>n;
121     Node* root=new Node(0);
122     for (int i=1;i<=n;i++)
123     {
124         cin>>opr>>x;
125         if (opr==1)
126         {
127             insert(root,x);
128         }
129         else if (opr==2)
130         {
131             if (!_find(root,x))
132                 cout<<"Not Found,Operation Failed"<<endl;
133             else
134                 remove(root,x);
135         }
136     }
137     cout<<"-----------------"<<endl;
138     cin>>m;
139     for (int i=1;i<=m;i++)
140     {
141         cin>>x;
142         if (_find(root,x))
143             cout<<"Found"<<endl;
144         else
145             cout<<"Not Found"<<endl;
146     }
147     cout<<"-----------------"<<endl;
148     cin>>m;
149     for (int i=1;i<=m;i++)
150     {
151         cin>>x;
152         int ans=kth(root,x);
153         cout<<x<<"th: "<<ans<<endl;
154     }
155 }