【牛客Wannafly挑战赛12】小H和圣诞树

题目

可以考虑边分治，对于某一种颜色，我们处理出分治边左右两边所有以这个颜色为端点的路径长度，之后随便拼一拼就好了

但是这样对于每一组询问都需要边分一遍，这样做复杂度是(O(nm+nlog n))的

还有一种更暴力的做法，就是枚举树上所有路径，这样就可以直接统计了，复杂度是(O(n^2))的

把这两个暴力结合一下，当一个分治块大小小于(sqrt{n})的时候，我们就直接跑第二种暴力，否则我们就跑边分治

跑第二种暴力的时候我们需要快速判断当前这个点对对应哪一个询问，于是需要二分一下，于是有一个(log)，所以我们把阈值设得比(sqrt{n})稍小一些，使得复杂度向边分治倾斜，这样就能跑过去了

代码

#include<bits/stdc++.h> 
#define re register
#define LL long long
#pragma GCC optimize(3)
#define max(a,b) ((a)>(b)?(a):(b))
inline int read() {
	char c=getchar();int x=0;while(c<'0'||c>'9') c=getchar();
	while(c>='0'&&c<='9') x=(x<<3)+(x<<1)+c-48,c=getchar();return x;
}
const int maxn=1e5+5;const int M=4e5+3;
int rn,n,num,rt,Mnow,S,B,tot,Q,sz;
struct E{int v,nxt,w;}e[M<<1];
std::vector<int> son[M],w[M];
int vis[M],head[M],sum[M],col[maxn];
LL Ans[maxn],tax[2][maxn],b[maxn];
int st[2][M],top[2],t[2][maxn],c[maxn],qx[maxn],qy[maxn];
inline int find(LL x) {
	int l=1,r=sz;
	while(l<=r) {
		int mid=l+r>>1;
		if(b[mid]==x) return mid;
		if(b[mid]<x) l=mid+1;else r=mid-1;
	}
	return 0;
}
inline void add(int x,int y,int w) {
	e[++num].v=y;e[num].nxt=head[x];head[x]=num;e[num].w=w;
}
void dfs1(int x,int fa) {
	for(re int i=head[x];i;i=e[i].nxt) {
		if(e[i].v==fa) continue;
		son[x].push_back(e[i].v);w[x].push_back(e[i].w);
		dfs1(e[i].v,x);
	}
}
void getrt(int x,int fa) {
	sum[x]=1;
	for(re int i=head[x];i;i=e[i].nxt) {
		if(e[i].v==fa||vis[i>>1]) continue;
		getrt(e[i].v,x);sum[x]+=sum[e[i].v];
		int now=max(sum[x],S-sum[x]);
		if(now<Mnow) Mnow=now,rt=i;
	}
}
void dfs2(int x,LL dis,int fa,int o) {
	if(x<=rn) tax[o][col[x]]+=dis,st[o][++top[o]]=col[x],t[o][col[x]]++;
	for(re int i=head[x];i;i=e[i].nxt) {
		if(e[i].v==fa||vis[i>>1]) continue;
		dfs2(e[i].v,dis+e[i].w,x,o);
	}
}
void calc(int x,int fa,int o,LL dis) {
	if(x<=rn&&col[x]>=o) {
		int to=c[find(1ll*o*rn+col[x])];
		if(to) Ans[to]+=dis;
	}
	for(re int i=head[x];i;i=e[i].nxt) {
		if(e[i].v==fa||vis[i>>1]) continue;
		calc(e[i].v,x,o,dis+e[i].w);
	}
}
void dfs(int x,int fa) {
	if(x<=rn) calc(x,0,col[x],0);
	for(re int i=head[x];i;i=e[i].nxt) {
		if(e[i].v==fa||vis[i>>1]) continue;
		dfs(e[i].v,x);
	}
}
void solve(int x,int s) {
	if(s<=B) {dfs(x,0);return;}
	Mnow=M,S=s,getrt(x,0);
	if(Mnow==M) return;vis[rt>>1]=1;
	dfs2(e[rt].v,e[rt].w,0,0);dfs2(e[rt^1].v,0,0,1);
	for(re int i=1;i<=Q;i++) {
		Ans[i]+=1ll*t[0][qx[i]]*tax[1][qy[i]]+1ll*t[1][qx[i]]*tax[0][qy[i]]
			+1ll*t[0][qy[i]]*tax[1][qx[i]]+1ll*t[1][qy[i]]*tax[0][qx[i]];
	} 
	while(top[0]) tax[0][st[0][top[0]]]=t[0][st[0][top[0]]]=0,top[0]--;
	while(top[1]) tax[1][st[1][top[1]]]=t[1][st[1][top[1]]]=0,top[1]--;
	int k=e[rt^1].v,h=S-sum[e[rt].v];
	solve(e[rt].v,sum[e[rt].v]);solve(k,h);
}
int main() {
	rn=n=read();
	for(re int i=1;i<=n;i++) col[i]=read();
	for(re int x,y,z,i=1;i<n;i++) 
		x=read(),y=read(),z=read(),add(x,y,z),add(y,x,z);
	dfs1(1,0);num=1;for(re int i=1;i<=n;i++) head[i]=0;int s[2];
	for(re int i=1;i<=n;i++) {
		int t=son[i].size();
		if(t<=2) {
			for(re int j=0;j<t;++j)
				add(i,son[i][j],w[i][j]),add(son[i][j],i,w[i][j]);
			continue;
		}
		s[0]=++n,s[1]=++n;
		add(i,s[0],0),add(s[0],i,0);add(i,s[1],0),add(s[1],i,0);
		for(re int j=0;j<t;++j)
			son[s[j&1]].push_back(son[i][j]),w[s[j&1]].push_back(w[i][j]);
	}
	Q=read();B=0.8*std::sqrt(rn);
	for(re int i=1;i<=Q;i++) {
		qx[i]=read(),qy[i]=read();
		if(qx[i]>qy[i]) std::swap(qx[i],qy[i]);
		b[i]=1ll*qx[i]*rn+qy[i];
	}
	std::sort(b+1,b+Q+1);sz=std::unique(b+1,b+Q+1)-b-1;
	for(re int i=Q;i;--i) c[find(1ll*qx[i]*rn+qy[i])]=i; 
	solve(1,n);
	for(re int i=1;i<=Q;i++) {
		int to=c[find(1ll*qx[i]*rn+qy[i])];
		if(to!=i) {printf("%lld
",Ans[to]);continue;}
		Ans[i]/=(qx[i]==qy[i]?2ll:1);
		printf("%lld
",Ans[i]);
	}
	return 0;
}