题解 P5838 【[USACO19DEC]Milk Visits(gold)】

这题一眼看出了可以用主席树维护暴力在线的做法，时间复杂度 (O(n log^2n))。方法类比宝石。

后来在想出了一个 (O(n log n)) 的做法，不需要任何高级数据结构。

这题用主席树的话，会发现所有查询都是已知的，不会有任何改变，且总查询次数是在 (O(m)) 级别的。

我们可以将所有需要的询问都先存下来，然后一遍 (dfs) 遍历解决所有询问，瓶颈在于求 LCA，时间复杂度 (O(n log n))。

#include<bits/stdc++.h>
#define pb push_back
#define F(i,a,b) for(int i=(a);i<=(b);++i)
#define DF(i,a,b) for(int i=(a);i>=(b);--i)
using namespace std;
inline int read(){char ch=getchar(); int w=1,c=0;
	for(;!isdigit(ch);ch=getchar()) if (ch=='-') w=-1;
	for(;isdigit(ch);ch=getchar()) c=(c<<1)+(c<<3)+(ch^48);
	return w*c;
}
const int N=1e5+10,M=1e5+10;
int n,m,t[N],dep[N],fa[25][N],a[M],b[M],c[M],k[M],pos[3][M],kind[N];
vector<int>v[N],qry[N],asw[N];
void dfs(int x,int fa){
	dep[x]=dep[fa]+1;::fa[0][x]=fa;
	F(i,1,20)::fa[i][x]=::fa[i-1][::fa[i-1][x]];
	for(int i:v[x])
		if(i!=fa)dfs(i,x);
}
int lca(int x,int y){
	if(dep[x]>dep[y])swap(x,y);
	DF(i,20,0)
		if(dep[fa[i][y]]>=dep[x])y=fa[i][y];
	if(x==y)return x;
	DF(i,20,0)
		if(fa[i][x]!=fa[i][y])x=fa[i][x],y=fa[i][y];
	return fa[0][x];
}
void dfs2(int x,int fa){
	kind[t[x]]++;
	for(int i:qry[x])asw[x].pb(kind[i]);
	for(int i:v[x])
		if(i!=fa)dfs2(i,x);
	kind[t[x]]--;
}
signed main(){
	n=read();m=read();
	F(i,1,n)t[i]=read();
	F(i,2,n){
		int x=read(),y=read();
		v[x].pb(y);
		v[y].pb(x);
	}
	dfs(1,0);
	F(i,1,m){
		a[i]=read(),b[i]=read(),c[i]=read();
		k[i]=lca(a[i],b[i]);
		pos[0][i]=qry[a[i]].size();
		qry[a[i]].pb(c[i]);
		pos[1][i]=qry[b[i]].size();
		qry[b[i]].pb(c[i]);
		pos[2][i]=qry[k[i]].size();
		qry[k[i]].pb(c[i]);
	}
	dfs2(1,0);
	F(i,1,m)
		putchar(asw[a[i]][pos[0][i]]+asw[b[i]][pos[1][i]]-asw[k[i]][pos[2][i]]*2+(t[k[i]]==c[i])?'1':'0');
	return 0;
}