树上游戏

这里是一个辣鸡的\(\Theta(n \log^2n)\)做法

首先我们套个点分治

然后问题就变成了统计跨过根的路径的贡献

先跑一遍以当前点为根的子树大小和DFS序

同时将根路径上首次出现该颜色的点存起来

然后根据颜色排序

对于不同颜色统计结果显然互不影响

那么对于同种颜色我们容斥一下就可以了

区间加法可以用DFS序上差分来代替

#include"cstdio"
#include"cstring"
#include"iostream"
#include"algorithm"
using namespace std;

const int MAXN=1e5+5;

int n,np,rt,sum,tp;
int h[MAXN];
int siz[MAXN],mx[MAXN],sz[MAXN],id[MAXN],rid[MAXN];
int clr[MAXN],bnk[MAXN],sm[MAXN];
long long v[MAXN];
bool vis[MAXN];
long long ans[MAXN];
struct rpg{
	int li,nx;
}a[MAXN<<1];
struct pt{
	int x,c,r;
}stk[MAXN];

int read()
{
	int x=0;char ch=getchar();
	while(ch<'0'||'9'<ch) ch=getchar();
	while('0'<=ch&&ch<='9') x=(x<<3)+(x<<1)+(ch^48),ch=getchar();
	return x;
}

void add(int ls,int nx)
{
	a[++np]=(rpg){h[ls],nx};h[ls]=np;
	a[++np]=(rpg){h[nx],ls};h[nx]=np;
}

void gtrt(int x,int fa)
{
	siz[x]=1;mx[x]=0;
	for(int i=h[x];i;i=a[i].li){
		if(a[i].nx==fa||vis[a[i].nx]) continue;
		gtrt(a[i].nx,x);
		siz[x]+=siz[a[i].nx];
		mx[x]=max(mx[x],siz[a[i].nx]);
	}mx[x]=max(mx[x],sum-siz[x]);
	if(mx[rt]>mx[x]) rt=x;
	return;
}

void gtsz(int x,int fa)
{
	sz[x]=1;id[x]=++id[0];rid[id[0]]=x;
	for(int i=h[x];i;i=a[i].li){
		if(a[i].nx==fa||vis[a[i].nx]) continue;
		gtsz(a[i].nx,x);
		sz[x]+=sz[a[i].nx];
	}return;
}

void dfs(int x,int fa,int rot)
{
	if(!bnk[clr[x]]) stk[++tp]=(pt){x,clr[x],rot};
	++bnk[clr[x]];
	for(int i=h[x];i;i=a[i].li){
		if(a[i].nx==fa||vis[a[i].nx]) continue;
		dfs(a[i].nx,x,rot);
	}--bnk[clr[x]];
	return;
}

bool cmp(pt a,pt b){return a.c<b.c;}

void calc(int x)
{
	id[0]=0;
	gtsz(x,x);bnk[clr[x]]=1;
   	ans[x]+=sz[x];//由根导致的起点为根的路径对根的贡献
	for(int i=h[x];i;i=a[i].li){
		if(vis[a[i].nx]) continue;
		dfs(a[i].nx,a[i].nx,a[i].nx);
		v[id[a[i].nx]]+=sz[x]-sz[a[i].nx];
		v[id[a[i].nx]+sz[a[i].nx]]-=sz[x]-sz[a[i].nx];
     		//由根导致的起点在该子树内跨过根的路径对起点的贡献
	}sort(stk+1,stk+tp+1,cmp);
	for(int l=1,r=0;l<=tp;l=r+1){
		int sum=0;
		while(r+1<=tp&&stk[r+1].c==stk[l].c) ++r,sm[stk[r].r]+=sz[stk[r].x],sum+=sz[stk[r].x];
		for(int i=l;i<=r;++i){
			v[id[stk[i].x]]-=sum-sm[stk[i].r];
			v[id[stk[i].x]+sz[stk[i].x]]+=sum-sm[stk[i].r];
			//减去跨过根的相同颜色导致多算的贡献
			v[1]+=sz[stk[i].x];
			v[id[stk[i].r]]-=sz[stk[i].x];
			v[id[stk[i].r]+sz[stk[i].r]]+=sz[stk[i].x];
			//统计终点在该子树内路径对起点的贡献
			v[id[stk[i].x]]+=sz[x]-sz[stk[i].r];
			v[id[stk[i].x]+sz[stk[i].x]]-=sz[x]-sz[stk[i].r];
        		//起点在该子树内对起点的贡献
		}for(int i=l;i<=r;++i) sm[stk[i].r]=0;
	}
	for(int i=1;i<=id[0];++i) v[i]+=v[i-1],ans[rid[i]]+=v[i];
	for(int i=1;i<=id[0]+1;++i) v[i]=0;
	bnk[clr[x]]=tp=0;
	return;
}

void DFZ(int x)
{
	vis[x]=1;calc(x);
	for(int i=h[x];i;i=a[i].li){
		if(vis[a[i].nx]) continue;
		sum=siz[a[i].nx];rt=0;
		gtrt(a[i].nx,a[i].nx);
		DFZ(rt);
	}return;
}

int main()
{
	n=read();
	for(int i=1;i<=n;++i) clr[i]=read();
	for(int i=1;i<n;++i){
		int x=read(),y=read();
		add(x,y);
	}sum=n;mx[rt]=n;
	gtrt(1,1);
	DFZ(rt);
	for(int i=1;i<=n;++i) printf("%lld\n",ans[i]);
	return 0;
}