【洛谷P7518】宝石

题目

题目链接：https://www.luogu.com.cn/problem/P7518
欧艾大陆上有 \(n\) 座城市，城市从 \(1 \sim n\) 编号，所有城市经由 \(n - 1\) 条无向道路互相连通，即 \(n\) 座城市与 \(n - 1\) 条道路构成了一棵树。
每座城市的集市上都会出售宝石，总共有 \(m\) 种不同的宝石，用 \(1 \sim m\) 编号。\(i\) 号城市的集市出售的是第 \(w_i\) 种宝石，一种宝石可能会在多座城市的集市出售。
K 神有一个宝石收集器。这个宝石收集器能按照顺序收集至多 \(c\) 颗宝石，其收集宝石的顺序为：\(P_1, P_2, \ldots , P_c\)。更具体地，收集器需要先放入第 \(P_1\) 种宝石，然后才能再放入第 \(P_2\) 种宝石，之后再能放入第 \(P_3\) 种宝石，以此类推。其中 \(P_1, P_2, \ldots , P_c\) 互不相等。
K 神到达一个城市后，如果该城市的集市上出售的宝石种类和当前收集器中需要放入的种类相同，则他可以在该城市的集市上购买一颗宝石并放入宝石收集器中；否则他只会路过该城市什么都不做。
现在 K 神给了你 \(q\) 次询问，每次给出起点 \(s_i\) 与终点 \(t_i\)，他想知道如果从 \(s_i\) 号城市出发，沿最短路线走到 \(t_i\) 号城市后，他的收集器中最多能收集到几个宝石？（在每次询问中，收集器内初始时没有任何宝石。起点与终点城市集市上的宝石可以尝试被收集）
\(n,q\leq 2\times 10^5,c\leq m\leq 5\times 10^4\)。

思路

下文 \(a[i]\) 表示点 \(i\) 的宝石编号，\(b[i]\) 表示按顺序收集的第 \(i\) 个宝石是哪一个。
因为保证了收集宝石的顺序中，每个宝石最多出现一次，所以我们可以直接记 \(\text{nxt}[i]\) 表示宝石 \(i\) 下一个需要收集的是哪一个宝石，\(\text{pre}[i]\) 表示 \(i\) 上一个收集的是哪一个宝石。
随便定一个节点为根，记 \(f[x][i]\) 表示点 \(x\) 开始收集宝石（\(x\) 必须选），一直往根节点走，收集到第 \(2^i\) 颗宝石时是在哪一个点。
只要我们对于每一个点 \(x\)，找到它到根的路径上第一个是 \(\text{nxt}[x]\) 的宝石，就可以倍增求出 \(f[x]\) 了。
建一棵主席树，\(x\) 为根的线段树内的一个叶子节点 \([i,i]\) 表示 \(x\) 到根的路径上，第一个是宝石 \(i\) 的节点是哪个。那么 \(x\) 的线段树只需要从 \(fa[x]\) 的线段树上修改 \(a[x]\) 这个位置即可。这样就可以在 \(O(n\log n)\) 的时间复杂度内求出 \(f[x][i]\)。
接下来考虑每一个询问 \(x,y\)。发现答案满足单调性，也就是如果可以拿到 \(k\) 个宝石，那么一定可以拿到 \(k-1\) 个宝石。所以考虑二分答案。
对于当前二分的答案 \(mid\)，利用主席树找到 \(x\) 到根的路径上第一个为 \(b[1]\) 的点，以及 \(y\) 到根的路径上第一个为 \(b[mid]\) 的点，设 \(\text{lca}(x,y)=p\)，我们可以利用 \(f\) 数组求出 \(x\) 到 \(p\) 的路径上能选多少个宝石，同理也可以求出在当前二分的 \(mid\) 下，\(y\) 到 \(p\) 的路径上能选出多少个宝石。注意因为实际上是从 \(p\) 走到 \(y\)，所以我们需要再预处理出一个 \(g[x][i]\) 表示 \(x\) 不断跳 \(\text{pre}[x]\) 跳 \(2^i\) 次到达的点。如果两条路径选出的宝石数量之和不小于 \(mid\)，则往更大的找，否则往更小的找。
时间复杂度 \(O(Q\log n\log m)\)。

代码

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

const int N=200010,LG=18;
int n,m,c,Q,tot,a[N],b[N],rt[N],nxt[N],pre[N],head[N],dep[N],pa[N][LG+1],f[N][LG+1],g[N][LG+1];

struct edge
{
	int next,to;
}e[N*2];

void add(int from,int to)
{
	e[++tot]=(edge){head[from],to};
	head[from]=tot;
}

struct SegTree
{
	int tot,lc[N*LG*4],rc[N*LG*4],val[N*LG*4];
	
	int update(int now,int l,int r,int k,int v)
	{
		int x=++tot;
		lc[x]=lc[now]; rc[x]=rc[now];
		if (l==r) { val[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);
		return x;
	}
	
	int query(int x,int l,int r,int k)
	{
		if (l==r) return val[x];
		int mid=(l+r)>>1;
		if (k<=mid) return query(lc[x],l,mid,k);
			else return query(rc[x],mid+1,r,k);
	}
}seg;

void dfs(int x,int fa)
{
	rt[x]=seg.update(rt[fa],1,m,a[x],x);
	dep[x]=dep[fa]+1; pa[x][0]=fa;
	if (nxt[a[x]]) f[x][0]=seg.query(rt[x],1,m,nxt[a[x]]);
	if (pre[a[x]]) g[x][0]=seg.query(rt[x],1,m,pre[a[x]]);
	for (int i=1;i<=LG;i++)
	{
		pa[x][i]=pa[pa[x][i-1]][i-1];
		f[x][i]=f[f[x][i-1]][i-1];
		g[x][i]=g[g[x][i-1]][i-1];
	}
	for (int i=head[x];~i;i=e[i].next)
	{
		int v=e[i].to;
		if (v!=fa) dfs(v,x);
	}
}

int lca(int x,int y)
{
	if (dep[x]<dep[y]) swap(x,y);
	for (int i=LG;i>=0;i--)
		if (dep[pa[x][i]]>=dep[y]) x=pa[x][i];
	if (x==y) return x; 
	for (int i=LG;i>=0;i--)
		if (pa[x][i]!=pa[y][i]) x=pa[x][i],y=pa[y][i];
	return pa[x][0];
}

int binary(int x,int y,bool typ)
{
	if (dep[x]<dep[y]) return 0;
	int res=1;
	for (int i=LG;i>=0;i--)
		if (!typ && dep[f[x][i]]>=dep[y])
			res+=(1<<i),x=f[x][i];
		else if (typ && dep[g[x][i]]>=dep[y])
			res+=(1<<i),x=g[x][i];
	return res;
}

int main()
{
	memset(head,-1,sizeof(head));
	scanf("%d%d%d",&n,&m,&c);
	for (int i=1;i<=c;i++)
	{
		scanf("%d",&b[i]);
		if (i>1) nxt[b[i-1]]=b[i],pre[b[i]]=b[i-1];
	}
	for (int i=1;i<=n;i++)
		scanf("%d",&a[i]);
	for (int i=1,x,y;i<n;i++)
	{
		scanf("%d%d",&x,&y);
		add(x,y); add(y,x);
	}
	dfs(1,0);
	scanf("%d",&Q);
	while (Q--)
	{
		int x,y;
		scanf("%d%d",&x,&y);
		int p=lca(x,y); x=seg.query(rt[x],1,m,b[1]);
		int cnt=binary(x,p,0),l=cnt+1,r=c,mid;
		while (l<=r)
		{
			mid=(l+r)>>1;
			int z=seg.query(rt[y],1,m,b[mid]);
			if (binary(z,p,1)+cnt>=mid) l=mid+1;
				else r=mid-1;
		}
		cout<<l-1<<"\n";
	}
	return 0;
}