权值线段树实现平衡树功能

最近想学习主席树 于是学了一个前置技能权值线段树 直接上代码qaq

#include<bits/stdc++.h>
#include<tr1/unordered_map>
#define For(i, a, b) for(register int i = a; i <= b; ++ i)

using namespace std;

const int maxn = 1e5 + 10;
unordered_map<int, int> id;
int m, cnt, rel[maxn];

struct node
{
    int k, x, id;
}Q[maxn];

bool cmp1(node A, node B)
{
    return A.x < B.x;
}

bool cmp2(node A, node B)
{   
    return A.id < B.id;
}

struct Segment_Tree
{
#define ls (bh << 1)
#define rs (bh << 1 | 1)
#define mid ((l + r) >> 1)
    int s[maxn << 2], Re;
    void pushup(int bh)
    {
        s[bh] = s[ls] + s[rs];
    }
    void update(int bh, int l, int r, int x, int z)
    {
        if(l == x && r == x)
            s[bh] += z;
        else
        {
            if(x <= mid)
                update(ls, l, mid, x, z);
            else
                update(rs, mid + 1, r, x, z);
            pushup(bh);
        }
    }
    int query(int bh, int l, int r, int x, int y)
    {
        if(x > y)
            return 0;
        int res = 0;
        if(x <= l && r <= y)
            res += s[bh];
        else
        {
            if(x <= mid)
                res += query(ls, l, mid, x, y);
            if(y > mid)
                res += query(rs, mid + 1, r, x, y);
        }
        return res;
    }
    int rank(int x)
    {
        return query(1, 1, cnt, 1, x - 1) + 1;
    }
    int Kth(int bh, int l, int r, int x)
    {
        if(l == r)
            return l; 
        if(s[ls] >= x)
            return Kth(ls, l, mid, x);
        return Kth(rs, mid + 1, r, x - s[ls]);
    }
    int find(int bh, int l, int r, int x)
    {
        if(l == r)
            return l;
        if(s[rs ^ x])
            return find(rs ^ x, x ? l : mid + 1, x ? mid : r, x);
        return find(ls ^ x, x ? mid + 1 : l, x ? r : mid, x);
    }
    int Pre(int bh, int l, int r, int x)
    {
        if(r < x)
        {
            if(s[bh])
                return find(bh, l, r, 0);
            return 0;
        }
        if(x > mid + 1 && s[rs] && (Re = Pre(rs, mid + 1, r, x)))
            return Re;
        return Pre(ls, l, mid, x);
    }
    int Sub(int bh, int l, int r, int x)
    {
        if(l > x)
        {
            if(s[bh])
                return find(bh, l, r, 1);
            return 0;
        }
        if(x < mid && s[ls] && (Re = Sub(ls, l, mid, x)))
            return Re;
        return Sub(rs, mid + 1, r, x);
    }
#undef mid
#undef ls
#undef rs
}T;

int main()
{
#ifndef ONLINE_JUDGE
    freopen("3369.in", "r", stdin);
    freopen("3369.out", "w", stdout);
#endif
    scanf("%d", &m);
    For(i, 1, m)
        scanf("%d%d", &Q[i].k, &Q[i].x), Q[i].id = i;
    sort(Q + 1, Q + m + 1, cmp1);
    For(i, 1, m)
        if(Q[i].k != 4 && !id[Q[i].x])
            id[Q[i].x] = ++ cnt, rel[cnt] = Q[i].x;
    sort(Q + 1, Q + m + 1, cmp2);
    For(i, 1, m)
    {
        if(Q[i].k == 1)
            T.update(1, 1, cnt, id[Q[i].x], 1);
        if(Q[i].k == 2)
            T.update(1, 1, cnt, id[Q[i].x], -1);
        if(Q[i].k == 3)
            printf("%d\n", T.rank(id[Q[i].x]));
        if(Q[i].k == 4)
            printf("%d\n", rel[T.Kth(1, 1, cnt, Q[i].x)]);
        if(Q[i].k == 5)
            printf("%d\n", rel[T.Pre(1, 1, cnt, id[Q[i].x])]);
        if(Q[i].k == 6)
            printf("%d\n", rel[T.Sub(1, 1, cnt, id[Q[i].x])]);
    }
    return 0;
}