[玄关] 初学 Treap 58pts WA7-12求条
查看原帖
[玄关] 初学 Treap 58pts WA7-12求条
612722
JustPureH2O楼主2024/7/15 21:01
#include <bits/stdc++.h>
#define N 100010
using namespace std;

struct Treap {
#define leftSubtree(idx) (tree[tree[idx].ls])
#define rightSubtree(idx) (tree[tree[idx].rs])
    int ls, rs;
    int dat;
    unsigned rnd;
    int cnt;
    int size;
} tree[N];

int idx = 0, root;
random_device rd;

int allocate(int dat, bool on_init = false) {
    idx++;
    tree[idx].dat = dat;
    tree[idx].rnd = (on_init) ? UINT_MAX : rd();
    tree[idx].cnt = (on_init) ? 0 : 1;
    tree[idx].size = (on_init) ? 0 : 1;
    return idx;
}

void build() {
    allocate(INT_MAX, true);
    allocate(-INT_MAX, true);
    tree[idx].ls = 2;
    root = 1;
}

void update(int idx) {
    tree[idx].size = (tree[idx].ls ? leftSubtree(idx).size : 0) + (tree[idx].rs ? rightSubtree(idx).size : 0) + tree[
                         idx].cnt;
}

void zig(int &p) {
    int q = tree[p].ls;
    tree[p].ls = tree[q].rs;
    tree[q].rs = p;
    p = q;
    update(tree[p].rs);
    update(p);
}

void zag(int &p) {
    int q = tree[p].rs;
    tree[p].rs = tree[q].ls;
    tree[q].ls = p;
    p = q;
    update(tree[p].ls);
    update(p);
}

void insert(int &idx, int dat) {
    if (!idx) idx = allocate(dat);
    else if (tree[idx].dat == dat) tree[idx].cnt++;
    else if (tree[idx].dat > dat) {
        insert(tree[idx].ls, dat);
        if (leftSubtree(idx).rnd > tree[idx].rnd) zig(idx);
    } else {
        insert(tree[idx].rs, dat);
        if (rightSubtree(idx).rnd > tree[idx].rnd) zag(idx);
    }
    update(idx);
}

void remove(int &idx, int dat) {
    if (!idx) return;
    if (tree[idx].dat == dat) {
        if (tree[idx].cnt > 1) tree[idx].cnt--;
        else {
            if (!tree[idx].ls) idx = tree[idx].rs;
            else if (!tree[idx].rs) idx = tree[idx].ls;
            else if (tree[idx].ls && tree[idx].rs) {
                if (leftSubtree(idx).rnd > tree[idx].rnd) {
                    zig(idx);
                    remove(tree[idx].rs, dat);
                } else if (rightSubtree(idx).rnd > tree[idx].rnd) {
                    zag(idx);
                    remove(tree[idx].ls, dat);
                }
            } else idx = 0;
        }
        update(idx);
        return;
    }
    if (tree[idx].dat > dat) remove(tree[idx].ls, dat);
    else remove(tree[idx].rs, dat);
    update(idx);
}

int getRankByData(int idx, int dat) {
    if (!idx) return 1;
    if (tree[idx].dat == dat) return (tree[idx].ls ? leftSubtree(idx).size : 0) + 1;
    if (tree[idx].dat > dat) return getRankByData(tree[idx].ls, dat);
    return getRankByData(tree[idx].rs, dat) + (tree[idx].ls ? leftSubtree(idx).size : 0) + tree[idx].cnt;
}

int getDataByRank(int idx, int rank) {
    if (!idx) return -INT_MAX;
    if (tree[idx].ls) {
        if (leftSubtree(idx).size >= rank) return getDataByRank(tree[idx].ls, rank);
        if (leftSubtree(idx).size + tree[idx].cnt >= rank) return tree[idx].dat;
    } else {
        if (idx > 2 && rank == 1) return tree[idx].dat;
    }
    return getDataByRank(tree[idx].rs, rank - (tree[idx].ls ? leftSubtree(idx).size : 0) - tree[idx].cnt);
}

int precursor(int idx, int dat) {
    if (!idx) return -INT_MAX;
    if (tree[idx].dat >= dat) return precursor(tree[idx].ls, dat);
    return max(tree[idx].dat, precursor(tree[idx].rs, dat));
}

int successor(int idx, int dat) {
    if (!idx) return INT_MAX;
    if (tree[idx].dat <= dat) return successor(tree[idx].rs, dat);
    return min(tree[idx].dat, successor(tree[idx].ls, dat));
}

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    int q;
    cin >> q;
    build();
    while (q--) {
        int op, x;
        cin >> op >> x;
        if (op == 1) insert(root, x);
        else if (op == 2) remove(root, x);
        else if (op == 3) cout << getRankByData(root, x) << endl;
        else if (op == 4) cout << getDataByRank(root, x) << endl;
        else if (op == 5) cout << precursor(root, x) << endl;
        else cout << successor(root, x) << endl;
    }
    return 0;
}

58pts记录

调了一整天,把同机房学过Treap的人都给整无语了……

2024/7/15 21:01
加载中...