(模板)treap

/*
key:关键字,sz:子树大小, cnt:节点重复个数
rand_val:随机权值,tot:节点总数
有标记的话需要在插入和旋转之前释放
*/

int ch[N][2];
int rand_val[N], key[N], sz[N], cnt[N];
int tot;

int creat(int val) {
    ++tot;
    ch[tot][0] = ch[tot][1] = 0;
    rand_val[tot] = rand();
    key[tot] = val;
    sz[tot] = cnt[tot] = 1;
    return tot;
}

int cmp(int x, int val) {
    if(key[x] == val) return -1;
    return val < key[x]? 0: 1;
}

void push_up(int x) {
    sz[x] = cnt[x] + sz[ch[x][0]] + sz[ch[x][1]];
}
void rot(int &x, int d) {
    int k = ch[x][!d];
    ch[x][!d] = ch[k][d];
    ch[k][d] = x;
    push_up(x);
    push_up(k);
    x = k;
}

void insert(int &x, int val) {
    if(x == 0) {
        x = creat(val);
        return;
    }
    int d = cmp(x, val);
    if(d == -1) {
        cnt[x]++, sz[x]++;
        return;
    }
    insert(ch[x][d], val);
    if(rand_val[ch[x][d]] > rand_val[x]) rot(x, d ^ 1);
    push_up(x);
}

void del(int &x, int val) {
    if(x == 0) return;
    int d = cmp(x, val);
    if(d == -1) {
        if(cnt[x] > 1) {
            cnt[x]--, sz[x]--;
            return;
        }
        if(!ch[x][0]) x = ch[x][1];
        else if(!ch[x][1]) x = ch[x][0];
        else {
            int d2 = rand_val[ch[x][0]] > rand_val[ch[x][1]]? 1: 0;
            rot(x, d2);
            del(ch[x][d2], val);
        }
    }
    else del(ch[x][d], val);
    if(x) push_up(x);
}

int kth(int x, int k) {
    if(!x) return -1;
    int tmp = sz[ch[x][0]];
    if(k <= tmp) return kth(ch[x][0], k);
    k -= tmp;
    if(k <= cnt[x]) return key[x];
    return kth(ch[x][1], k - cnt[x]);
}

你可能感兴趣的:((模板)treap)