LOJ 104. 普通平衡树

#include <bits/stdc++.h>
using namespace std;
const int N = 1e5 + 10;
int n, opt, x, val[N], fa[N], sze[N], sum[N], lc[N], rc[N], T, rt;
void pushup(int k) {
    sze[k] = sze[lc[k]] + sze[rc[k]] + sum[k];
}
void rot(int x) {
    int y = fa[x], z = fa[y];
    int b = (x == lc[y]) ? rc[x] : lc[x];
    fa[x] = z, fa[y] = x;
    if (b) fa[b] = y;
    if (z) (lc[z] == y ? lc[z] : rc[z]) = x;
    if (x == lc[y]) rc[x] = y, lc[y] = b;
    else lc[x] = y, rc[y] = b;
    pushup(y);
}
bool Wrt(int x) {
    return rc[fa[x]] == x;
}    
void Splay(int x, int tar) {
    while (fa[x] != tar) {
        if (fa[fa[x]] != tar) 
            Wrt(x) == Wrt(fa[x]) ? rot(fa[x]) : rot(x);
        rot(x);
    }
    if (!tar) rt = x;
    pushup(x);
}
void ins(int k) {//插入元素k 
    int x = rt, y = 0, d;
    while (x) {
        ++ sze[y = x];
        if (k == val[x]) {sum[x] ++; return ;}
        if (k < val[x]) x = lc[x], d = 0;
        else x = rc[x], d = 1;
    }
    x = ++ T; fa[x] = y; sze[x] = sum[x] = 1; val[x] = k;
    if (y) (d == 0 ? lc[y] : rc[y]) = x;
    Splay(x, 0);
}
int Rank(int k, bool f) { //按数值找   0：位置   1：排名
     int x = rt, ret = 0;
    while (x) {
        if (val[x] == k) return f ? ret + sze[lc[x]] + 1 : x;
        if (val[x] > k) x = lc[x];
        else ret += sze[lc[x]] + sum[x], x = rc[x]; 
    }    
    return f ? ret : 0;
}
int find(int k) {//按排名找 
    int x = rt;
    while (x) {
        if (sze[lc[x]] >= k) x = lc[x];
        else if (sze[lc[x]] + sum[x] >= k) return val[x];
        else k -= (sze[lc[x]] + sum[x]), x = rc[x];
    }
    return val[x];
}
void del(int k) {//删除元素k所在节点 
    int x = Rank(k, 0);
    Splay(x, 0);
    if (sum[x] > 1) {
        -- sze[x]; -- sum[x]; return ;
    }
    fa[lc[x]] = fa[rc[x]] = 0;
    if (!lc[x] || !rc[x]) rt = lc[x] + rc[x];
    else {
        int w = lc[x];
        while(rc[w]) w = rc[w];
        Splay(w, 0);
        rc[w] = rc[x]; fa[rc[x]] = w;
        pushup(w);
    }
}
int pre(int k) {//求前驱 
    int x = rt, ret = 0;
    while (x) {
        if (val[x] < k) ret = val[x], x = rc[x];
        else x = lc[x];
    }
    return ret;
}
int aft(int k) {//求后继 
    int x = rt, ret = 0;
    while (x) {
        if (val[x] > k) ret = val[x], x = lc[x];
        else x = rc[x];
    }
    return ret;
}
signed main() {        
    scanf("%d", &n);
    while (n --) {
        scanf("%d %d", &opt, &x);
        if (opt == 1) ins(x);
        if (opt == 2) del(x);
        if (opt == 3) cout << Rank(x,1) << '\n';
        if (opt == 4) cout << find(x) << '\n';
        if (opt == 5) cout << pre(x) << '\n';
        if (opt == 6) cout << aft(x) << '\n';
    }    
    return 0;
}

 

---

如果您觉得阅读本文对您有帮助，请点一下“推荐”按钮，您的“推荐”将是我最大的写作动力！欢迎各位转载，但是未经作者本人同意，转载文章之后必须在文章页面明显位置给出作者和原文连接，否则保留追究法律责任的权利。