标题说平衡树,那么应该AVL,红黑树都能过,但是这次做这题主要是学习Treap,所以花了几天搞出了这题。其他方法以后再说吧

  • Treap
    #include <bits/stdc++.h>
    using namespace std;
    const int MAXN = 100010;
    const int INF = 0x3f3f3f3f;
    struct Treap {
        int son[2];
        int val, rand;
        int cnt, size;
    } node[MAXN];
    int root, tot;
    void push_up(int i) {
        int ls = node[i].son[0];
        int rs = node[i].son[1];
        node[i].size = node[ls].size + node[rs].size + node[i].cnt;
    }
    void rotate(int& r, int index) {
        int k = node[r].son[index ^ 1];
        node[r].son[index ^ 1] = node[k].son[index];
        node[k].son[index] = r;
        push_up(r); push_up(k);
        r = k;
    }
    void insert(int& r, int v) {
        if (r == 0) {
            r = ++ tot;
            node[r].val = v;
            node[r].rand = rand();
            node[r].size = node[r].cnt = 1;
            return;
        }
        if (node[r].val == v) {
            node[r].size ++;
            node[r].cnt ++;
            return;
        }
        int index = v > node[r].val;
        insert(node[r].son[index], v);
        if (node[node[r].son[index]].rand < node[r].rand) rotate(r, index ^ 1);
        push_up(r);
    }
    void erase(int& r, int v) {
        if (r == 0) return;
        if (v < node[r].val) erase(node[r].son[0], v);
        else if (v > node[r].val) erase(node[r].son[1], v);
        else {
            if (node[r].son[0] == 0 && node[r].son[1] == 0) {
                node[r].cnt --;
                node[r].size --;
                if (node[r].cnt == 0) r = 0;
            } else if (node[r].son[0] != 0 && node[r].son[1] == 0) {
                rotate(r, 1);
                erase(node[r].son[1], v);
            } else if (node[r].son[0] == 0 && node[r].son[1] != 0) {
                rotate(r, 0);
                erase(node[r].son[0], v);
            } else {
                int ls = node[r].son[0];
                int rs = node[r].son[1];
                int index = node[ls].rand < node[rs].rand;
                rotate(r, index);
                erase(node[r].son[index], v);
            }
        }
        push_up(r);
    }
    int get_rank(int r, int v) {
        int size = node[node[r].son[0]].size;
        int cnt = node[r].cnt;
        if (node[r].val == v) return size + 1;
        if (node[r].val > v) return get_rank(node[r].son[0], v);
        if (node[r].val < v) return size + cnt + get_rank(node[r].son[1], v);
    }
    int get_val(int r, int rk) {
        int size = node[node[r].son[0]].size;
        int cnt = node[r].cnt;
        if (rk < size + 1) return get_val(node[r].son[0], rk);
        if (rk >= size + 1 && rk <= size + cnt) return node[r].val;
        if (rk > size + cnt) return get_val(node[r].son[1], rk - size - cnt);
    }
    int get_pre(int r, int v) {
        if (r == 0) return -INF;
        if (node[r].val >= v) return get_pre(node[r].son[0], v);
        else return max(node[r].val, get_pre(node[r].son[1], v));
    }
    int get_suc(int r, int v) {
        if (r == 0) return INF;
        if (node[r].val <= v) return get_suc(node[r].son[1], v);
        else return min(node[r].val, get_suc(node[r].son[0], v));
    }
    int main() {
        srand(time(NULL));
        int n; scanf("%d", &n);
        for (int i = 1; i <= n; i++) {
            int opt, x;
            scanf("%d%d", &opt, &x);
            if (opt == 1) insert(root, x);
            if (opt == 2) erase(root, x);
            if (opt == 3) printf("%d\n", get_rank(root, x));
            if (opt == 4) printf("%d\n", get_val(root, x));
            if (opt == 5) printf("%d\n", get_pre(root, x));
            if (opt == 6) printf("%d\n", get_suc(root, x));
        }
        return 0;
    }

    学到的第二个用随机数的算法,所以这个Treap的高度也是有的随机的。但是起码不太可能退化成链状。

01-03 09:54
查看更多