Splay 求调
查看原帖
Splay 求调
570503
zaochen楼主2023/10/25 22:15

大红大紫,但是AC了一半。
不知道为啥在哪个OJ都只能AC一半。

#include <bits/stdc++.h>
using namespace std;
// #define int long long

const int N = 1e5 + 5, inf = 1 << 30;

struct node {
    int ch[2], fa;    // 左右孩子节点,父亲节点的编号
    int val, cnt, sz; // 当前节点的值,当前节点有多少个,当前子树的大小
    #define lc (t[p].ch[0])
    #define rc (t[p].ch[1])
} t[N];
int tot, root;
int a[N], n, q;

int new_node(int val) { // 新建节点
    int x = ++tot;
    t[x].val = val, t[x].cnt = t[x].sz = 1;
    return x;
}
void connect(int p, int fa, int k) { // 把节点 p 设置为节点 fa 的左儿子或右儿子
    t[p].fa = fa, t[fa].ch[k] = p;
}
int get(int p) { // 判断节点 p 是其父亲的左子节点或右子节点
    return p == t[t[p].fa].ch[1];
}
void push_up(int p) { // 向上传递信息
    t[p].sz = t[lc].sz + t[rc].sz + t[p].cnt;
}

void rotate(int p) { // 单旋
    int fa = t[p].fa, ffa = t[fa].fa, k = get(p);
    connect(p, ffa, get(fa));
    connect(t[p].ch[k ^ 1], fa, k);
    connect(fa, p, k ^ 1);
    push_up(fa);
    push_up(p);
}

void splay(int p, int k = 0) { // 把节点 p 伸展到 k 的子节点,k = 0 时表示伸展到根节点
    if (!k)
        root = p;

    while (t[p].fa != k) {
        if (t[t[p].fa].fa != k)
            get(p) ^ get(t[p].fa) ? rotate(p) : rotate(t[p].fa);

        rotate(p);
    }
}

void insert(int& p, int val, int fa) { // 插入节点
    if (!p) {
        p = new_node(val);
        t[p].fa = fa;
        splay(p);
        return;
    }
    else if (val < t[p].val)
        insert(lc, val, p);
    else if (val > t[p].val)
        insert(rc, val, p);
    else
        t[p].cnt++;
}
void build() { // 新建一颗 Splay 树
    root = new_node(inf);
    insert(root, -inf, 0);
}

int get_prev(int p, int val) { // 找前驱节点
    if (!p)
        return -1;

    if (t[p].val >= val)
        return get_prev(lc, val);

    int r = get_prev(rc, val);

    if (r == -1 || t[p].val > t[r].val) {
        if (p == root)
            splay(p);

        return p;
    }
    else {
        if (p == root)
            splay(r);

        return r;
    }

    // return max(t[p].val, get_prev(rc, val));
}
int get_next(int p, int val) { // 找后继节点z
    if (!p)
        return -1;

    if (t[p].val <= val)
        return get_next(rc, val);

    int l = get_next(lc, val);

    if (l == -1 || t[p].val < t[l].val) {
        if (p == root)
            splay(p);

        return p;
    }
    else {
        if (p == root)
            splay(l);

        return l;
    }

    // return min(t[p].val, get_next(lc, val));
}

int kth(int p, int k) { // 找排名为 k 的值
    if (k <= t[lc].sz)
        return kth(lc, k);
    else if (k <= t[lc].sz + t[p].cnt) {
        splay(p);
        return t[p].val;
    }

    return kth(rc, k - t[lc].sz - t[p].cnt);
}
int rk(int p, int val) { // 查询 val 的排名
    if (!p)
        return 1;

    if (t[p].val == val) {
        return t[lc].sz + 1;
    }

    if (t[p].val > val)
        return rk(lc, val);

    return t[lc].sz + t[p].cnt + rk(rc, val);
}

void remove(int p, int val) { // 删除一个值为 val 的结点
    if (val < t[p].val)
        remove(lc, val);
    else if (val > t[p].val)
        remove(rc, val);
    else {
        splay(p);

        if (t[p].cnt > 1) {
            t[p].cnt--;
            push_up(p);
            return;
        }

        int r = get_next(t[p].ch[1], p);
        splay(r, p);
        connect(lc, r, 0);
        // t[t[root].ch[1]].ch[0] = lc;

        root = r;
        t[r].fa = 0;
        push_up(root);
    }
}

void output(int p) { // 中序遍历
    #ifdef LUOTIANYI
    if (lc)
        output(lc);

    if (p)
        cout << t[p].val << ' ' << t[p].cnt << endl;

    if (rc)
        output(rc);

    #endif
}

signed main() {
    ios::sync_with_stdio(false);
    #ifdef LUOTIANYI
    clock_t t0 = clock();
    freopen("data.in", "r", stdin);
    freopen("data.out", "w", stdout);
    #endif

    // Don't stop. Don't hide. Follow the light, and you'll find tomorrow.
    cin >> n;
    build();

    for (int i = 1; i <= n; i++) {
        int op, x;
        cin >> op >> x;

        if (op == 1)
            insert(root, x, 0);
        else if (op == 2)
            remove(root, x);
        else if (op == 3)
            cout << (rk(root, x) - 1) << '\n';
        else if (op == 4)
            cout << kth(root, x + 1) << '\n';
        else if (op == 5)
            cout << t[get_prev(root, x)].val << '\n';
        else
            cout << t[get_next(root, x)].val << '\n';
    }

    #ifdef LUOTIANYI
    cerr << "Time used:" << clock() - t0 << "ms" << endl;
    #endif
    return 0;
}
2023/10/25 22:15
加载中...