【洛谷7518】[省选联考 2021 A/B 卷] 宝石（树上倍增+并查集）

点此看题面

• 有一棵\(n\)个点的树，每棵树上有一种颜色的宝石。
• 有一个宝石收集器，能按给定顺序收集\(c\)颗宝石（保证这\(c\)颗宝石颜色各不相同）。
• 每次询问从\(x\)走到\(y\)，若当前点的宝石与当前收集宝石颜色相同则收集，求收集的宝石颗数。
• \(n,q\le2\times10^5\)

重新标色

我们给宝石重新标颜色，让要收集的第\(i\)颗宝石颜色变成\(i\)，而无法收集的宝石颜色变成\(0\)。

因为题目保证这\(c\)颗宝石颜色各不相同，应该是很好实现的。

这样一来就方便许多。

向上：树上倍增

我们把一条询问链拆成从\(x\)向上到\(LCA\)的子节点和从\(LCA\)向下到\(y\)的两部分。

向上的部分可以倍增预处理一个数组\(f_{x,i}\)表示已收集到第\(a_x\)颗宝石，从\(x\)向上接着按顺序收集到第\(a_x+2^i\)颗宝石需要到达哪个节点，并预处理出\(g_x\)表示从\(x\)向上深度最大的第\(1\)颗宝石所在节点。

要预处理这两个数组，我们只需\(dfs\)一遍，维护从根节点到当前点的这条链上每种宝石深度最大的点\(h_i\)，则初始化\(f_{x,0}=h_{a_x+1}\)（接下来的倍增和普通倍增完全没区别），\(g_x=h_1\)。

那么，我们首先从\(x\)跳到\(g_x\)，因为必须要从第\(1\)颗宝石开始，然后倍增上跳，满足深度大于等于\(dep_{LCA(x,y)}+1\)即可。

向下：并查集

求出了向上的答案，我们把询问编号和向上答案绑成一个二元组扔到\(LCA\)的\(vector\)中，并把询问编号扔到另一个询问节点\(y\)的另一\(vector\)中表示需要在\(y\)点上询问。

我们记录\(c+1\)个根节点\(rt_{0\sim c}\)，这样就可以维护出已收集到每种宝石的询问集合。

处理到一个点时，我们首先取出\(vector\)中的二元组，把询问对应的节点加入向上答案对应的集合中。

然后，我们可以收集当前点的宝石\(a_i\)，因此把\(a_{i}-1\)合并到\(a_i\)中，然后把 \(rt_{a_i-1}\) 修改成\(0\)。

接着就是处理当前点的询问，找到它对应节点所在连通块即可（为此我们还需要对每个根节点记录它是哪个连通块的根节点）。

递归处理完子树后，离开这个点时我们还需要撤销这个点的贡献，因此使用按秩合并可撤销并查集。

代码：\(O(nlogn)\)

#include<bits/stdc++.h>
#define Tp template<typename Ty>
#define Ts template<typename Ty,typename... Ar>
#define Reg register
#define RI Reg int
#define Con const
#define CI Con int&
#define I inline
#define W while
#define N 200000
#define LN 18
#define add(x,y) (e[++ee].nxt=lnk[x],e[lnk[x]=ee].to=y)
using namespace std;
int n,m,c,a[N+5],p[N+5],ans[N+5],ee,lnk[N+5];struct edge {int to,nxt;}e[2*N+5];
namespace FastIO
{
	#define FS 100000
	#define tc() (FA==FB&&(FB=(FA=FI)+fread(FI,1,FS,stdin),FA==FB)?EOF:*FA++)
	#define pc(c) (FC==FE&&(clear(),0),*FC++=c)
	int OT;char oc,FI[FS],FO[FS],OS[FS],*FA=FI,*FB=FI,*FC=FO,*FE=FO+FS;
	I void clear() {fwrite(FO,1,FC-FO,stdout),FC=FO;}
	Tp I void read(Ty& x) {x=0;W(!isdigit(oc=tc()));W(x=(x<<3)+(x<<1)+(oc&15),isdigit(oc=tc()));}
	Ts I void read(Ty& x,Ar&... y) {read(x),read(y...);}
	Tp I void write(Ty x) {W(OS[++OT]=x%10+48,x/=10);W(OT) pc(OS[OT--]);pc(' ');}
	Tp I void writeln(Ty x) {W(OS[++OT]=x%10+48,x/=10);W(OT) pc(OS[OT--]);pc('\n');}
}using namespace FastIO;
int h[N+5],g[N+5],dep[N+5],f[N+5][LN+5],fa[N+5][LN+5];I void Init(CI x)//预处理
{
	RI i;for(i=1;i<=LN;++i) fa[x][i]=fa[fa[x][i-1]][i-1];RI o=h[a[x]];h[a[x]]=x;//更新h数组
	for(g[x]=h[1],f[x][0]=h[a[x]+1],i=1;i<=LN;++i) f[x][i]=f[f[x][i-1]][i-1];//利用h预处理f和g
	for(RI i=lnk[x];i;i=e[i].nxt) e[i].to^fa[x][0]&&
		(dep[e[i].to]=dep[fa[e[i].to][0]=x]+1,Init(e[i].to),0);h[a[x]]=o;//撤销影响，还原h数组
}
I int Q(RI x,CI d)//倍增
{
	if(dep[x]<d) return 0;for(RI i=LN;~i;--i) dep[f[x][i]]>=d&&(x=f[x][i]);return a[x];
}
I int LCA(RI x,RI y)//LCA
{
	RI i;for(dep[x]<dep[y]&&(swap(x,y),0),i=0;dep[x]^dep[y];++i) (dep[x]^dep[y])>>i&1&&(x=fa[x][i]);
	if(x==y) return x;for(i=LN;~i;--i) fa[x][i]^fa[y][i]&&(x=fa[x][i],y=fa[y][i]);return fa[x][0];
}
namespace U//按秩合并可撤销并查集
{
	int f[N+5],g[N+5],T,Sx[2*N+5],Sy[2*N+5],Gx[2*N+5],Gy[2*N+5];
	I int fa(CI x) {return f[x]?fa(f[x]):x;}
	I void Union(int& x,RI y)//合并，由于必然是合并两个无父节点的节点，不需要getfa
	{
		Sx[++T]=x,Sy[T]=y,Gx[T]=g[x],Gy[T]=g[y];if(!x||!y) return (void)(x|=y);
		g[x]<g[y]&&(swap(x,y),0),g[f[y]=x]==g[y]&&++g[x];
	}
	I void Back(int& x) {g[x=Sx[T]]=Gx[T],g[Sy[T]]=Gy[T],f[Sx[T]]=f[Sy[T]]=0,--T;}//撤销上次合并
}
struct node {int p,v;};vector<node> s[N+5];vector<int> q[N+5];
vector<node>::iterator st;vector<int>::iterator qt;
int w[N+5],rt[N+5];I void Solve(CI x)//求解
{
	for(st=s[x].begin();st!=s[x].end();++st) U::Union(rt[st->v],st->p),w[rt[st->v]]=st->v;//把询问加到向上答案对应集合中
	RI o;a[x]&&(U::Union(rt[a[x]],rt[a[x]-1]),o=rt[a[x]-1],rt[a[x]-1]=0,rt[a[x]]&&(w[rt[a[x]]]=a[x]));//把a[x]-1合并到a[x]
	for(qt=q[x].begin();qt!=q[x].end();++qt) ans[*qt]=w[U::fa(*qt)];//处理询问
	for(RI i=lnk[x];i;i=e[i].nxt) e[i].to^fa[x][0]&&(Solve(e[i].to),0);//递归处理子树
	a[x]&&(U::Back(rt[a[x]]),rt[a[x]]&&(w[rt[a[x]]]=a[x]),(rt[a[x]-1]=o)&&(w[rt[a[x]-1]]=a[x]-1));//撤销a[x]-1到a[x]的合并
	for(st=s[x].end();st!=s[x].begin();) --st,U::Back(rt[st->v]),w[rt[st->v]]=st->v;//撤销询问的加入
}
int main()
{
	RI Qt,i,x,y,z;for(read(n,m,c),i=1;i<=c;++i) read(x),p[x]=i;
	for(i=1;i<=n;++i) read(a[i]),a[i]=p[a[i]];for(i=1;i^n;++i) read(x,y),add(x,y),add(y,x);dep[1]=1,Init(1);
	for(read(Qt),i=1;i<=Qt;++i) read(x,y),z=LCA(x,y),s[z].push_back((node){i,Q(g[x],dep[z]+1)}),q[y].push_back(i);//拆成向上和向下两部分
	for(Solve(1),i=1;i<=Qt;++i) writeln(ans[i]);return clear(),0;
}