python/TangDou/AcWing/SegmentTree/P5494.cpp

#include <cstdio>
#include <cstring>
#include <algorithm>
#include <iostream>
#include <cmath>

using namespace std;
typedef long long LL;

//重定义 tr[u].l 和 tr[u].r 快速录入用
#define ls tr[u].lson
#define rs tr[u].rson

const int N = 200010;
int n, m;
int a[N]; //记录原始数组，其实也可以不用保存，直接在build时录入到叶子也可以

int root[N], idx; //分裂出的第几个线段树,idx是序号维护器

struct Node {
    int lson, rson; //动态开点,记录左右儿子的节点编号
                    //注意：这里与普通线段树不同，没有记录当前节点的管辖范围！！！！！
                    //黄海尝试了记录[l,r]的方法，结果因为结构体内增加了两个属性，同时数组的上限非常大，直接MLE了好多测试点。
    LL val;         //每一个节点保存的值大小是管辖区间中的个数
} tr[N << 6];
int cnt; //记录tr数组中的可以放入节点的位置

//权值线段树合并,维护区间元素个数
void pushup(int u) {
    tr[u].val = tr[ls].val + tr[rs].val;
}

/**
 * @brief 构建线段树
 *
 * @param l 想要创建的节点，是管理哪个范围的，这与普通线段树相似，只不过普通线段树可以事先准备好节点号，这个需要现用现创建，
 * 并且需要返回节点号，让父节点记录和父节点的关联关系。
 * @param r
 * @return int
 */
int build(int l, int r) {
    //由于构建之前，空间是未分配状态，需要为新构建的节点分配节点号，节点号就是cnt
    int u = ++cnt;
    //现在有了可用的节点号了
    if (l == r) { //如果是叶子节点
        //接下来一行n个整数，表示 1∼n 这些数在a[N]中出现的 **次数**
        //其实7就放在[7,7]的位置上，val就记录7这个数字出现的次数
        tr[u].val = a[l]; //个数！！！
        return u;
    }
    int mid = (l + r) >> 1;

    //由于是动态开点，无法像普通线段树一样，采用预先创建的办法，普通办法只要知道父节点，u<<1 和 u<<1|1就是左右儿子节点号
    //而在动态开点的模板中，需要递归创建并返回新建节点编号，由并记录到tr[u].l和tr[u].r上，以实现父子节点的关联对应
    ls = build(l, mid), rs = build(mid + 1, r);
    //更新父节点统计信息
    pushup(u);
    return u;
}
/**
 * @brief 线段树分裂
 *
 * @param k 以k为根的线段树
 * @param l 左边界
 * @param r 右边界
 * @param ql 要分裂出的左边界
 * @param qr 要分裂出的右边界
 * @return int
 */
int spilt(int k, int l, int r, int ql, int qr) {
    int u = ++cnt;                               //分裂后的根节点编号u
    if (l == ql && r == qr) {                    //完全命中区间
        tr[u] = tr[k];                           //将tr[k]抄到tr[u]
        tr[k].val = tr[k].lson = tr[k].rson = 0; //销毁tr[k]
        return u;                                //返回新分裂后的节点编号u
    }
    int mid = (l + r) >> 1;
    if (qr <= mid) //分裂左儿子
        ls = spilt(tr[k].lson, l, mid, ql, qr);
    else if (ql > mid) //分裂右儿子
        rs = spilt(tr[k].rson, mid + 1, r, ql, qr);
    else {
        ls = spilt(tr[k].lson, l, mid, ql, mid);
        rs = spilt(tr[k].rson, mid + 1, r, mid + 1, qr);
    }
    //递归更新分裂后u和k的父节点信息
    pushup(u), pushup(k);
    return u;
}

//合并线段树
void merge(int &x, int y) {
    if (!(x && y))
        x |= y;
    else {
        tr[x].val += tr[y].val;
        merge(tr[x].lson, tr[y].lson);
        merge(tr[x].rson, tr[y].rson);
    }
}

/**
 * @brief 在 u 这个可重集中加入 x 个数字 q
 *
 * @param u 以u为根
 * @param l 左边界
 * @param r 右边界
 * @param q 数字p
 * @param x 增加x个
 */
void insert(int u, int l, int r, int q, int x) {
    if (l == r) {
        tr[u].val += x; //权值线段树个数增加x个
        return;
    }
    int mid = (l + r) >> 1;
    if (q <= mid) {
        if (ls == 0) ls = ++cnt;  //左儿子不存在，则创建之
        insert(ls, l, mid, q, x); //向左儿子递归插入
    } else {
        if (rs == 0) rs = ++cnt;      //右儿子不存在，则创建之
        insert(rs, mid + 1, r, q, x); //向右儿子递归插入
    }
    //向父节点更新统计信息
    pushup(u);
}
/**
 * @brief 在以u节点为根的线段树中，管辖范围是[l,r],查询区间[ql,qr]内数字的总个数
 *        查询可重集 p中大于等于 x 且小于等于 y 的值的个数。
 * @param u 根节点
 * @param l 左边界
 * @param r 右边界
 * @param ql 查询左边界
 * @param qr 查询右边界
 * @return LL 有多少个
 */
LL query(int u, int l, int r, int ql, int qr) {
    if (l == ql && r == qr) return tr[u].val;
    int mid = (l + r) >> 1;
    if (qr <= mid)
        return query(ls, l, mid, ql, qr);
    else if (ql > mid)
        return query(rs, mid + 1, r, ql, qr);
    else
        return query(ls, l, mid, ql, mid) + query(rs, mid + 1, r, mid + 1, qr);
}

//查询第k小的数
int kth(int u, int l, int r, int k) {
    if (l == r) return l;
    int mid = (l + r) >> 1;
    if (tr[ls].val >= k)
        return kth(ls, l, mid, k);
    return kth(rs, mid + 1, r, k - tr[ls].val);
}

int main() {
    //加快读入
    ios::sync_with_stdio(false), cin.tie(0);

    cin >> n >> m;
    //范围[1~n],创建第一个线段树，返回值是根节点号，记录到root[1]中
    for (int i = 1; i <= n; i++) cin >> a[i];
    root[++idx] = build(1, n);

    int op, p, x, y;
    for (int i = 1; i <= m; i++) {
        cin >> op >> p;
        if (op == 0) { //将可重集p中大于等于x且小于等于y的值放入一个新的可重集中
            //新可重集编号为从2开始的正整数，是上一次产生的新可重集的编号+1。
            cin >> x >> y;
            root[++idx] = spilt(root[p], 1, n, x, y);
        } else if (op == 1) { //合并线段树
            cin >> x;
            merge(root[p], root[x]); //将根为x的线段树合并到根为p的线段树中去,x线段树以后就不用了
        } else if (op == 2) {        //在p这个可重集中加入x个数字q
            cin >> x >> y;
            insert(root[p], 1, n, y, x);
        } else if (op == 3) { //区间查询
            cin >> x >> y;
            printf("%lld\n", query(root[p], 1, n, x, y));
        } else {
            int k;
            cin >> k; //在线段树上二分，如果左孩子的元素个数大于等于k，说明第k小在左子树内；
            //否则，在右子树内。
            if (query(root[p], 1, n, 1, n) < k)
                printf("-1\n");
            else
                printf("%d\n", kth(root[p], 1, n, k));
        }
    }
    return 0;
}