P5838 [USACO19DEC]Milk Visits G 【树上莫队】

P5838 [USACO19DEC]Milk Visits G

给定一颗树，每个点有点权，\(m\) 次询问，每次求 \(u\) 到 \(v\) 之间的路径是否出现过权值为 \(w\) 的点。

树上莫队板子题，我们需要一个 dfs 序来帮助进行莫队操作，dfs 序记录一个 \(in\) 一个 \(out\)，所以我们有 \(2 \times n\) 的 dfs 序，每次读入时先将 \(in\) 的大小选出 \(l,r\) ，求出 \(lca(l,r)\) 。

然后分两种情况讨论：如果 \(l\) 是 \(r\) 的 \(lca\) （包含 \(l=r\))，则标记 \(in[l],in[r],lca=-1\) ；不然则认为是一条普通路径，标记 \(out[l],in[r],lca=lca(l,r)\) 。这样操作保证 \([l,r]\) 中间的点只会被遍历一次，如果 \(lca != -1\) ，则需要指针移动完毕时加入 \(lca\) 再记录答案，然后删除 \(lca\) 的贡献。

因为 dfs 序的关系，不同与普通序列，我们额外需要为莫队指针使用一个 \(vis\) 数组记录当前的点是否出现，出现则加入，否则删除，每次 \(vis[u] xor = 1\)。区别与普通莫队的 add，del 。树上莫队只需要一个 \(work\) 就可以同时实现两个操作了。

Code.

#include<bits/stdc++.h>
using namespace std;
const int N=1e5+10;
int n,m,h[N],ne[N<<1],e[N<<1],idx,len,cnt,st[N],vis[N],sz[N],w[N],in[N],out[N],top[N],son[N],fa[N],dep[N],pl[N<<1],ans[N];
inline int get(int x) {return x/len;}
struct node
{
	int l,r,w,lca,id;
	bool operator <(const node &o) const {
		int i=get(l),j=get(o.l);
		if(i == j) return r < o.r; return i < j;
	} 
} q[N];
void add(int u,int v) {ne[++idx]=h[u],e[idx]=v,h[u]=idx;}
void dfs1(int u,int father,int depth)
{
	in[u]=++cnt; pl[cnt]=u; fa[u]=father; dep[u]=depth; sz[u]=1;
	for(int i=h[u];~i;i=ne[i])
	{
		int j=e[i]; if(j == father) continue ;
		dfs1(j,u,depth+1); sz[u]+=sz[j];
		if(sz[son[u]] < sz[j]) son[u]=j;
	}
	out[u]=++cnt; pl[cnt]=u;
}
void dfs2(int u,int t)
{
	top[u]=t; if(! son[u]) return ; dfs2(son[u],t);
	for(int i=h[u];~i;i=ne[i])
	{
		int j=e[i];
		if(j == fa[u] || j == son[u]) continue ;
		dfs2(j,j);
	}
}
inline int lca(int u,int v)
{
	while(top[u] != top[v])
	{
		if(dep[top[u]] < dep[top[v]]) swap(u,v);
		u=fa[top[u]];
	}
	return dep[u] < dep[v] ? u : v;
}
inline void work(int x) {if(vis[pl[x]] == 0) st[w[pl[x]]]++,vis[pl[x]]=1; else st[w[pl[x]]]--,vis[pl[x]]=0; }
int main()
{
	memset(h,-1,sizeof h); scanf("%d%d",&n,&m);
	for(int i=1;i<=n;i++) scanf("%d",&w[i]);
	for(int i=1,u,v;i<n;i++) {scanf("%d%d",&u,&v); add(u,v); add(v,u);}
	len=(int)sqrt(2*n)+1; dfs1(1,0,1); dfs2(1,1);
	for(int i=1,u,v;i<=m;i++)
	{
		scanf("%d%d%d",&u,&v,&q[i].w); q[i].id=i; if(in[u] > in[v]) swap(u,v); 
		if(u == v) {q[i].l=in[u],q[i].r=in[v],q[i].lca=-1; continue ;} 
		int yl=lca(u,v);
		if(yl == u) q[i].l=in[u],q[i].r=in[v],q[i].lca=-1;
		else q[i].l=out[u],q[i].r=in[v],q[i].lca=yl;
	}
	stable_sort(q+1,q+1+m);
	int l=1,r=0;
	for(int i=1;i<=m;i++)
	{
		int l1=q[i].l,r1=q[i].r;
		while(l < l1) work(l++);
		while(l > l1) work(--l);
		while(r < r1) work(++r);
		while(r > r1) work(r--);
		if(q[i].lca != -1) work(in[q[i].lca]);
		ans[q[i].id]=st[q[i].w];
		if(q[i].lca != -1) work(in[q[i].lca]);
	}
	for(int i=1;i<=m;i++) if(ans[i] == 0) cout<<0; else cout<<1;
	return 0;
}