模板 - 线段树套平衡树

一种可以在区间上找名次的数据结构。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;

namespace Treap {
#define ls ch[id][0]
#define rs ch[id][1]
    const int INF = 2147483647;
    const int N = 5e4 + 5;
    const int MAXN = 25 * N;
    //每个元素会被覆盖若干次线段树有log层，每层的平衡树的长度和都是n，按道理说是精确的nlogn，但这里还是开尽可能大
    int ch[MAXN][2], dat[MAXN];

    int val[MAXN];
    int cnt[MAXN];
    int siz[MAXN];

    int tot;

    inline void Init() {
        tot = 0;
    }

    inline int NewNode(int v, int num) {
        int id = ++tot;
        ls = rs = 0;
        dat[id] = rand();
        val[id] = v;
        cnt[id] = num;
        siz[id] = num;
        return id;
    }

    inline void PushUp(int id) {
        siz[id] = siz[ls] + siz[rs] + cnt[id];
    }

    inline void Rotate(int &id, int d) {
        int temp = ch[id][d ^ 1];
        ch[id][d ^ 1] = ch[temp][d];
        ch[temp][d] = id;
        id = temp;
        PushUp(ch[id][d]);
        PushUp(id);
    }

    //插入num个v
    inline void Insert(int &id, int v, int num) {
        if(!id)
            id = NewNode(v, num);
        else {
            if(v == val[id])
                cnt[id] += num;
            else {
                int d = val[id] > v ? 0 : 1;
                Insert(ch[id][d], v, num);
                if(dat[id] < dat[ch[id][d]])
                    Rotate(id, d ^ 1);
            }
            PushUp(id);
        }
    }

    //删除至多num个v
    void Remove(int &id, int v, int num) {
        if(!id)
            return;
        else {
            if(v == val[id]) {
                if(cnt[id] > num) {
                    cnt[id] -= num;
                    PushUp(id);
                } else if(ls || rs) {
                    if(!rs || dat[ls] > dat[rs])
                        Rotate(id, 1), Remove(rs, v, num);
                    else
                        Rotate(id, 0), Remove(ls, v, num);
                    PushUp(id);
                } else
                    id = 0;
            } else {
                val[id] > v ? Remove(ls, v, num) : Remove(rs, v, num);
                PushUp(id);
            }
        }
    }

    //查询严格<v的数的个数（和普通平衡树不一样
    int GetRank(int id, int v) {
        int res = 0;
        while(id) {
            if(val[id] > v)
                id = ls;
            else if(val[id] == v) {
                res += siz[ls];
                break;
            } else {
                res += siz[ls] + cnt[id];
                id = rs;
            }
        }
        return res;
    }

    //查询树中的<=x的数的个数有至少rk个的最小的x
    int GetValue(int id, int rk) {
        int res = INF;
        while(id) {
            if(siz[ls] >= rk)
                id = ls;
            else if(siz[ls] + cnt[id] >= rk) {
                res = val[id];
                break;
            } else {
                rk -= siz[ls] + cnt[id];
                id = rs;
            }
        }
        return res;
    }

    //查询v的前驱的值（<v的第一个节点的值），不存在前驱返回负无穷
    int GetPrev(int id, int v) {
        int res = -INF;
        while(id) {
            if(val[id] < v)
                res = val[id], id = rs;
            else
                id = ls;
        }
        return res;
    }

    //查询v的后继的值（>v的第一个节点的值），不存在后继返回无穷
    int GetNext(int id, int v) {
        int res = INF;
        while(id) {
            if(val[id] > v)
                res = val[id], id = ls;
            else
                id = rs;
        }
        return res;
    }
#undef ls
#undef rs
}

namespace SegmentTree {
#define ls (p<<1)
#define rs (p<<1|1)
    const int INF = 2147483647;
    const int MAXN = 5e4 + 5;
    int n, m, a[MAXN];
    struct SegmentTreeNode {
        int l, r;
        int root;
    } st[MAXN * 4];
    //线段树的大小是精确的4倍

    void Build(int p, int l, int r) {
        st[p].l = l, st[p].r = r;
        for (int i = l; i <= r ; ++i)
            Treap::Insert(st[p].root, a[i], 1);
        if(l == r)
            return;
        int mid = l + r >> 1;
        Build(ls, l, mid);
        Build(rs, mid + 1, r);
    }

    void Update(int p, int pos, int v) {
        Treap::Remove(st[p].root, a[pos], 1);
        Treap::Insert(st[p].root, v, 1);
        if (st[p].l == st[p].r)
            return;
        int mid = st[p].l + st[p].r >> 1;
        if(pos <= mid)
            Update(ls, pos, v);
        else
            Update(rs, pos, v);
    }

    //查询严格<v的数的个数（和普通平衡树不一样
    int GetRank(int p, int l, int r, int v) {
        if (st[p].l > r || st[p].r < l)
            return 0;
        if (st[p].l >= l && st[p].r <= r)
            return Treap::GetRank(st[p].root, v);
        else
            //很明显是满足结合律的
            return GetRank(ls, l, r, v) + GetRank(rs, l, r, v) ;
    }

    //查询区间中的<=x的数的个数有至少rk个的最小的x
    int GetValue(int l, int r, int k) {
        int L = 0, R = 1e8;
        while(1) {
            int mid = L + R  >> 1;
            if(L == mid)
                return L;
            //小于k，也有可能mid是刚刚好的，比如mid有连续的一段
            if(GetRank(1, l, r, mid) < k)
                L = mid;
            else
                R = mid;
        }
    }

    int GetPrev(int p, int l, int r, int k) {
        if (st[p].l > r || st[p].r < l)
            return -INF;
        if (st[p].l >= l && st[p].r <= r)
            return Treap::GetPrev(st[p].root, k);
        else
            //很明显是满足结合律的
            return max(GetPrev(ls, l, r, k), GetPrev(rs, l, r, k));
    }

    int GetNext(int p, int l, int r, int k) {
        if (st[p].l > r || st[p].r < l)
            return INF;
        if (st[p].l >= l && st[p].r <= r)
            return Treap::GetNext(st[p].root, k);
        else
            //很明显是满足结合律的
            return min(GetNext(ls, l, r, k), GetNext(rs, l, r, k));
    }
#undef ls
#undef rs
}

int main() {
    int n, m;
    scanf("%d%d", &n, &m);
    SegmentTree::n = n;
    for (int i = 1; i <= n ; ++i)
        scanf("%d", &SegmentTree::a[i]);

    SegmentTree::Build(1, 1, n);
    for (int i = 1; i <= m; ++i) {
        int opt, l, k, r, pos;
        scanf("%d", &opt);
        switch(opt) {
        case 1:
            scanf("%d%d%d", &l, &r, &k);
            printf("%d\n", SegmentTree::GetRank(1, l, r, k) + 1);
            break;
        case 2:
            scanf("%d%d%d", &l, &r, &k);
            printf("%d\n", SegmentTree::GetValue(l, r, k));
            break;
        case 3:
            scanf("%d%d", &pos, &k);
            SegmentTree::Update(1, pos, k);
            SegmentTree::a[pos] = k;
            break;
        case 4:
            scanf("%d%d%d", &l, &r, &k);
            printf("%d\n", SegmentTree::GetPrev(1, l, r, k));
            break;
        case 5:
            scanf("%d%d%d", &l, &r, &k);
            printf("%d\n", SegmentTree::GetNext(1, l, r, k));
            break;
        }
    }
    return 0;
}

转载于:https://www.cnblogs.com/Inko/p/11414166.html