【洛谷P5494】【模板】线段树分裂

题目

题目链接：https://www.luogu.com.cn/problem/P5494
给出一个可重集 \(a\)（编号为 \(1\)），它支持以下操作：
0 p x y：将可重集 \(p\) 中大于等于 \(x\) 且小于等于 \(y\) 的值放入一个新的可重集中（新可重集编号为从 \(2\) 开始的正整数，是上一次产生的新可重集的编号+1）。
1 p t：将可重集 \(t\) 中的数放入可重集 \(p\)，且清空可重集 \(t\)（数据保证在此后的操作中不会出现可重集 \(t\)）。
2 p x q：在 \(p\) 这个可重集中加入 \(x\) 个数字 \(q\)。
3 p x y：查询可重集 \(p\) 中大于等于 \(x\) 且小于等于 \(y\) 的值的个数。
4 p k：查询在 \(p\) 这个可重集中第 \(k\) 小的数，不存在时输出 -1。

\(n,m\leq 2\times 10^5\)。

思路

操作 \(1,2,3,4\) 可以用权值线段树和线段树合并解决。至于操作 \(1\) 线段树合并的复杂度，每次可以看作是 \(O(\text{节点数较小值})\) 的，所以一般的线段树合并是 \(O(n\log n)\)。在本题中，操作 \(0,2\) 单次加入的节点数均为 \(O(\log n)\)，所以复杂度没有问题。
考虑操作 \(0\)，也就是把权值线段树的一个区间 \([l,r]\) 分为 \([l,k]\) 和 \([k+1,r]\) 扔到两棵权值线段树上。类似 FHQ Treap 的 split，对于当前区间 \([l,r]\)，如果 \(k\leq mid\)，则左边的权值线段树的这一层节点为空，右边的直接设为 \(x\)；反之同理。这样就保证了每次只会找 \(O(\log n)\) 个节点。
而合并 \([0,x),(y,n]\) 区间的时候，由于这两棵线段树没有交，单次复杂度也是 \(O(\log n)\) 的。
时间复杂度 \(O(m\log n)\)。

代码

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

const int N=200010,LG=18;
int n,m,Q,rt[N];

struct SegTree
{
	int tot,lc[N*LG*4],rc[N*LG*4];
	ll cnt[N*LG*4];
	
	void pushup(int x)
	{
		cnt[x]=cnt[lc[x]]+cnt[rc[x]];
	}
	
	int update(int x,int l,int r,int k,ll v)
	{
		if (!x) x=++tot;
		if (l==r) { cnt[x]+=v; return x; }
		int mid=(l+r)>>1;
		if (k<=mid) lc[x]=update(lc[x],l,mid,k,v);
			else rc[x]=update(rc[x],mid+1,r,k,v);
		pushup(x);
		return x;
	}
	
	void split(int x,int l,int r,int v,int &p,int &q)
	{
		if (!x) return;
		if (l==r) { p=x; q=0; return; }
		int mid=(l+r)>>1;
		if (v<=mid)
			q=x,p=++tot,split(lc[x],l,mid,v,lc[p],lc[q]);
		else
			p=x,q=++tot,split(rc[x],mid+1,r,v,rc[p],rc[q]);
		pushup(p); pushup(q);
	}
	
	int merge(int x,int y,int l,int r)
	{
		if (!x || !y) return x|y;
		if (l==r) { cnt[x]+=cnt[y]; return x; }
		int mid=(l+r)>>1;
		lc[x]=merge(lc[x],lc[y],l,mid);
		rc[x]=merge(rc[x],rc[y],mid+1,r);
		pushup(x);
		return x;
	}
	
	ll query1(int x,int l,int r,int ql,int qr)
	{
		if (!x) return 0;
		if (ql<=l && qr>=r) return cnt[x];
		int mid=(l+r)>>1; ll res=0;
		if (ql<=mid) res+=query1(lc[x],l,mid,ql,qr);
		if (qr>mid) res+=query1(rc[x],mid+1,r,ql,qr);
		return res;
	}
	
	int query2(int x,int l,int r,ll k)
	{
		if (l==r) return l;
		int mid=(l+r)>>1;
		if (cnt[lc[x]]>=k) return query2(lc[x],l,mid,k);
			else return query2(rc[x],mid+1,r,k-cnt[lc[x]]);
	}
}seg;

int main()
{
//	freopen("P5494_1.in","r",stdin);
//	freopen("ans.out","w",stdout);
	scanf("%d%d",&n,&Q);
	for (int i=1,x;i<=n;i++)
	{
		scanf("%d",&x);
		rt[1]=seg.update(rt[1],1,n,i,x);
	}
	m=1;
	while (Q--)
	{
		int opt,x,y,z;
		scanf("%d",&opt);
		if (opt==0)
		{
			int p=0,q=0;
			scanf("%d%d%d",&x,&y,&z);
			if (y>1) seg.split(rt[x],1,n,y-1,p,rt[x]);
			seg.split(rt[x],1,n,z,rt[++m],q);
			rt[x]=seg.merge(p,q,1,n);
		}
		if (opt==1)
		{
			scanf("%d%d",&x,&y);
			rt[x]=seg.merge(rt[x],rt[y],1,n);
		}
		if (opt==2)
		{
			scanf("%d%d%d",&x,&y,&z);
			rt[x]=seg.update(rt[x],1,n,z,y);
		}
		if (opt==3)
		{
			scanf("%d%d%d",&x,&y,&z);
			cout<<seg.query1(rt[x],1,n,y,z)<<"\n";
		}
		if (opt==4)
		{
			ll k;
			scanf("%d",&x); scanf("%lld",&k);
			if (seg.cnt[rt[x]]<k) cout<<"-1\n";
				else cout<<seg.query2(rt[x],1,n,k)<<"\n";
		}
	}
	return 0;
}