SPOJ 3273 - Order statistic set Treap

點擊打開鏈接

題意：
集合S支持一下四種操作：
INSERT(S,x) : 如果S中沒有x，則插入x
DELETE(S,x): 如果S中有x，則刪除x
K-TH(S): 輸出S中第K小的數
COUNT(S,x): 統計S中小於x的數有多少個

一共有Q（1 ≤ Q ≤ 200000）次操作。

Treap模板。。

#include
#include
#include
const int inf=0x3f3f3f3f;

struct Node {
  Node *ch[2]; 
  int r, v, s;
  Node(int v):v(v) { ch[0] = ch[1] = NULL; r = rand(); s = 1; }
  
  int cmp(int x) const {
    if (x == v) return -1;
    return (x < v) ? 0 : 1;
  }

  void maintain() {
    s = 1;
    if(ch[0] != NULL) s += ch[0]->s;
    if(ch[1] != NULL) s += ch[1]->s;
  }
};

void rotate(Node* &o, int d) {
  Node* k = o->ch[d^1]; o->ch[d^1] = k->ch[d]; k->ch[d] = o;
  o->maintain(); k->maintain(); o = k;
}

void insert(Node* &o, int x) {
  if(o == NULL) o = new Node(x);
  else{
    /*int d = (x < o->v ? 0 : 1); // 允許插入相同結點*/
    int d =o->cmp(x); if(d==-1) return ; //不允許插入相同結點*/
    insert(o->ch[d], x);
    if(o->ch[d]->r > o->r) rotate(o, d^1);
    else o->maintain();
  }
}

void remove(Node* &o, int x) {
  int d = o->cmp(x);
  if(d == -1) {
    Node* u = o;
    if(o->ch[0] != NULL && o->ch[1] != NULL) {
      int d2 = (o->ch[0]->r > o->ch[1]->r) ? 1 : 0;
      rotate(o, d2); remove(o->ch[d2], x);
    }else{
      if(o->ch[0] == NULL) o = o->ch[1]; else o = o->ch[0];
      delete u;
    }
  }else remove(o->ch[d], x);
  if(o != NULL) o->maintain();
}

int find(Node* o, int x){
    while(o != NULL){
        int d = o->cmp(x);
        if(d == -1) return 1;
        o = o->ch[d];
    }
    return 0;
}

int kth(Node* o,int k){///kth small
    while(o != NULL){
        int s = (o->ch[0] == NULL) ? 0 : o->ch[0]->s;
        if(k==s+1) return o->v;
        else if(k<=s) o=o->ch[0];
        else{
            o=o->ch[1]; k=k-s-1;
        }
    }
    return -inf;
}

int rank(Node* o, int x){
    int ret = 0;
    while(o != NULL){
        int s = (o->ch[0] == NULL) ? 0 : o->ch[0]->s;
        if(o->v < x){
            ret += s + 1;
            o = o->ch[1];
        }else{
            o = o->ch[0];
        }
    }
    return ret;
}


int main()
{
    int q;
    Node* root=NULL;
    //freopen("in.txt","r",stdin);
    scanf("%d",&q);
    while(q--)
    {
        char cmd[3];int nb;
        scanf("%s%d",cmd,&nb);
        if(cmd[0]=='I'&&find(root,nb)==0)
        {
            insert(root,nb);
        }
        else if(cmd[0]=='D'&&find(root,nb)==1)
        {
            remove(root,nb);
        }
        else if(cmd[0]=='K')
        {
            int temp=kth(root,nb);
            if(temp == -inf) puts("invalid");
            else printf("%d\n",temp);
        }
        else if(cmd[0]=='C')
        {
            printf("%d\n",rank(root,nb));
        }
    }
    return 0;
}