【模板】Splay

Luogu P3369 【模板】普通平衡树

#include <iostream>
#include <stdio.h>
#include <string.h>
#define MAX_N 100005

using namespace std;

int rt = 0, tot = 0;
int par[MAX_N];
int ch[MAX_N][2];
int w[MAX_N];
int cnt[MAX_N];
int siz[MAX_N];

void pushup(int x) {
    siz[x] = siz[ch[x][0]] + siz[ch[x][1]] + cnt[x];
}

int get(int x) {
    return ch[par[x]][1] == x;
}

void rot(int x) {
    int y = par[x], z = par[y], c = get(x), d = get(y);
    ch[y][c] = ch[x][c ^ 1], par[ch[y][c]] = y;
    ch[x][c ^ 1] = y, par[y] = x;
    par[x] = z;
    if (z) ch[z][d] = x;
    pushup(y), pushup(x);
}

void splay(int x) {
    for (int y; y = par[x]; rot(x)) {
        if (par[y]) {
            rot(get(x) == get(y) ? y : x);
        }
    }
    rt = x;
}

void ins(int v) {
    if (!rt) {
        w[++tot] = v, cnt[tot]++, rt = tot;
        pushup(rt);
        return;
    }
    int cur = rt, p = 0;
    while (true) {
        if (w[cur] == v) {
            cnt[cur]++;
            pushup(cur), pushup(p);
            splay(cur);
            return;
        }
        p = cur, cur = ch[cur][v > w[cur]];
        if (!cur) {
            w[++tot] = v, cnt[tot]++, par[tot] = p, ch[p][v > w[p]] = tot;
            pushup(tot), pushup(p);
            splay(tot);
            return;
        }
    }
}

int rk(int v) {
    int res = 0, cur = rt;
    while (true) {
        if (v < w[cur]) {
            cur = ch[cur][0];
        } else {
            res += siz[ch[cur][0]];
            if (w[cur] == v) {
                splay(cur);
                return res + 1;
            }
            res += cnt[cur];
            cur = ch[cur][1];
        }
    }
}

int kth(int k) {
    int cur = rt;
    while (true) {
        if (k <= siz[ch[cur][0]]) {
            cur = ch[cur][0];
        } else {
            k -= siz[ch[cur][0]] + cnt[cur];
            if (k <= 0) {
                splay(cur);
                return w[cur];
            }
            cur = ch[cur][1];
        }
    }
}

int pre() {
    int cur = ch[rt][0];
    while (ch[cur][1]) cur = ch[cur][1];
    splay(cur);
    return cur;
}

int suc() {
    int cur = ch[rt][1];
    while (ch[cur][0]) cur = ch[cur][0];
    splay(cur);
    return cur;
}

void del(int v) {
    rk(v);
    if (cnt[rt] > 1) {
        cnt[rt]--;
        pushup(rt);
        return;
    }
    if (!ch[rt][0] && !ch[rt][1]) {
        rt = 0;
        return;
    }
    if (!ch[rt][0] || !ch[rt][1]) {
        rt = ch[rt][0] + ch[rt][1];
        par[rt] = 0;
        return;
    }
    int cur = rt, x = pre();
    ch[x][1] = ch[cur][1], par[ch[x][1]] = x;
    pushup(x);
}

int main() {
    ios::sync_with_stdio(false);
    int n, op, x;
    cin >> n;
    while (n--) {
        cin >> op >> x;
        if (op == 1) {
            ins(x);
        } else if (op == 2) {
            del(x);
        } else if (op == 3) {
            cout << rk(x) << endl;
        } else if (op == 4) {
            cout << kth(x) << endl;
        } else if (op == 5) {
            ins(x);
            cout << w[pre()] << endl;
            del(x);
        } else {
            ins(x);
            cout << w[suc()] << endl;
            del(x);
        }
    }
}
点赞

发表评论

电子邮件地址不会被公开。必填项已用 * 标注