模板 - 数据结构 - Treap

Posted 2021-03-21 kisekipurin2019

还有人把Treap叫做树堆的，但是常用名还是叫做Treap的比较多。

不进行任何封装的，带求和操作的，一个节点存放多个元素的最普通的Treap。

#include<bits/stdc++.h>
using namespace std;
typedef long long ll;
#define ls ch[id][0]
#define rs ch[id][1]

const int INF = 1e9;

const int MAXN = 1000000 + 5;
int ch[MAXN][2], dat[MAXN];

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

int tot, root;

inline void Init() {
    tot = 0;
    root = 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;
    sum[id] = 1ll * num * v;
    return id;
}

inline void PushUp(int id) {
    siz[id] = siz[ls] + siz[rs] + cnt[id];
    sum[id] = sum[ls] + sum[rs] + 1ll * cnt[id] * val[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的排名，排名定义为<v的数的个数+1。
int GetRank(int id, int v) {
    int res = 1;
    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;
}

//查询排名为rk的数，rk必须是正整数，rk过大返回无穷
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;
}

//查询小于等于v的数的和
ll GetSumValue(int id, int v) {
    ll res = 0;
    while(id) {
        if(val[id] > v)
            id = ls;
        else if(val[id] == v) {
            res += sum[ls] + 1ll * cnt[id] * val[id];
            break;
        } else {
            res += sum[ls] + 1ll * cnt[id] * val[id];
            id = rs;
        }
    }
    return res;
}

//查询前rk个数的和，rk必须是正整数
ll GetSumRank(int id, int rk) {
    ll res = 0;
    while(id) {
        if(siz[ls] >= rk)
            id = ls;
        else if(siz[ls] + cnt[id] >= rk) {
            res += sum[ls] + 1ll * (rk - siz[ls]) * val[id];
            break;
        } else {
            res += sum[ls] + 1ll * cnt[id] * val[id];
            rk -= siz[ls] + cnt[id];
            id = rs;
        }
    }
    return res;
}

封装了val的，速度略微下降，因为是键值对所以求和类的函数变得没什么意义。

struct TreapNode {
    int val1, val2;
    TreapNode() {}
    TreapNode(int val1, int val2): val1(val1), val2(val2) {}

    bool operator<(const TreapNode& tn)const {
        return val1 < tn.val1;
    }
    bool operator<=(const TreapNode& tn)const {
        return val1 <= tn.val1;
    }
    bool operator==(const TreapNode& tn)const {
        return val1 == tn.val1;
    }
    bool operator>=(const TreapNode& tn)const {
        return val1 >= tn.val1;
    }
    bool operator>(const TreapNode& tn)const {
        return val1 > tn.val1;
    }
} TNINF(INF, INF);

/*#define TreapNode pii
TreapNode TNINF(INF, INF);*/

struct Treap {
#define ls ch[id][0]
#define rs ch[id][1]
    static const int MAXN = 200000;
    int ch[MAXN + 5][2], dat[MAXN + 5];

    TreapNode val[MAXN + 5];

    int cnt[MAXN + 5];
    int siz[MAXN + 5];

    int tot, root;

    void Init() {
        tot = 0;
        root = 0;
    }

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

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

    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
    void _Insert(int &id, TreapNode 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, TreapNode 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的排名，排名定义为<v的数的个数+1。
    int _GetRank(int id, TreapNode v) {
        int res = 1;
        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;
    }

    //查询排名为rk的数，rk必须是正整数，rk过大返回无穷
    TreapNode _GetValue(int id, int rk) {
        TreapNode res = TNINF;
        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;
    }

    int Size() {
        return siz[root];
    }
    void Insert(TreapNode v, int num = 1) {
        _Insert(root, v, num);
    }
    void Remove(TreapNode v, int num = INF) {
        _Remove(root, v, num);
    }
    int GetRank(TreapNode v) {
        return _GetRank(root, v);
    }
    TreapNode GetValue(int rk) {
        return _GetValue(root, rk);
    }
#undef ls
#undef rs
}