「LG2664 树上游戏」

题目

这真是一道神仙的一批的题目

定义\(s(i,j)\)表示从点\(i\)到点\(j\)经过的颜色数量

设

\[sum_i=\sum_{j=1}^ns(i,j) \]

求出所有的\(sum_i\)

考虑点分治

对于一个点我们用两种方式来统计其答案

1. 这个点作为分治重心时，分值区域内所有点到这个点贡献

2. 这个点不是分治重心的时候，当前分治区域内其他子树到这个点的贡献

第一种贡献我们很好统计，点分治的时候把所有子树遍历一遍就好了

第二种就需要转换一下思路了，我们不能直接求\(s(i,j)\)了，我们应该求某一种颜色一共被数了多少次

我们开一个桶\(tax\)，\(tax[i]\)表示\(i\)这种颜色控制的大小一共是多少，也就是这个颜色会被多少个终点数到，我们可以通过提前遍历好所有子树得到这个信息

每次进入一棵子树的时候，提前减掉这个子树的贡献，之后进入子树\(dfs\)就好了，如果一旦出现一种新颜色，显然这种颜色会被当前分治区域内所有点数上，更改一下贡献即可

代码

#include<algorithm>
#include<iostream>
#include<cstring>
#include<cstdio>
#include<vector>
#define maxn 100005
#define re register
#define inf 99999999
#define LL long long
#define max(a,b) ((a)>(b)?(a):(b))
#define min(a,b) ((a)<(b)?(a):(b))
inline int read() {
	int x=0;char c=getchar();while(c<'0'||c>'9') c=getchar();
	while(c>='0'&&c<='9') x=(x<<3)+(x<<1)+c-48,c=getchar();return x;
}
struct E{int v,nxt;}e[maxn<<1];
int col[maxn],head[maxn],vis[maxn],sum[maxn],mx[maxn];
int f[maxn],tax[maxn],d[maxn],st[maxn],tmp[maxn];
int num,n,m,now,S,rt,top;
LL ans,Ans[maxn],res;
std::vector<int> v[maxn],c[maxn];
inline void add(int x,int y) {e[++num].v=y;e[num].nxt=head[x];head[x]=num;}
void getroot(int x,int fa) {
	sum[x]=1,mx[x]=0;
	for(re int i=head[x];i;i=e[i].nxt) {
		if(vis[e[i].v]||e[i].v==fa) continue;
		getroot(e[i].v,x);
		sum[x]+=sum[e[i].v],mx[x]=max(mx[x],sum[e[i].v]);
	}
	mx[x]=max(mx[x],S-sum[x]);
	if(mx[x]<now) now=mx[x],rt=x;
}
void getdis(int x,int fa,int now,int t) {
	if(!f[col[x]]) now++;
	if(!tmp[col[x]]) st[++top]=col[x];
	tmp[col[x]]=1;sum[x]=1;f[col[x]]++;Ans[t]+=now;
	for(re int i=head[x];i;i=e[i].nxt) {
		if(vis[e[i].v]||e[i].v==fa) continue;
		getdis(e[i].v,x,now,t);sum[x]+=sum[e[i].v];
	}
	if(f[col[x]]==1) d[col[x]]+=sum[x];
	f[col[x]]--;
}
void find(int x,int fa) {
	if(!f[col[x]]) ans-=tax[col[x]],ans+=res;
	Ans[x]+=ans;f[col[x]]++;
	for(re int i=head[x];i;i=e[i].nxt) {
		if(vis[e[i].v]||e[i].v==fa) continue;
		find(e[i].v,x);
	} 
	if(f[col[x]]==1) ans-=res,ans+=tax[col[x]];
	f[col[x]]--;
}
void dfs(int x) {
	vis[x]=1;ans=0;f[col[x]]=1;
	for(re int i=head[x];i;i=e[i].nxt) {
		if(vis[e[i].v]) continue;
		top=0;getdis(e[i].v,0,1,x);
		for(re int j=1;j<=top;j++) 
			if(st[j]!=col[x]) v[e[i].v].push_back(d[st[j]]),c[e[i].v].push_back(st[j]);
		for(re int j=1;j<=top;j++) 
			if(st[j]!=col[x]) tax[st[j]]+=d[st[j]],ans+=d[st[j]];
		for(re int j=1;j<=top;j++) tmp[st[j]]=0,d[st[j]]=0;
	}
	f[col[x]]=0;
	ans+=S,tax[col[x]]=S;
	for(re int i=head[x];i;i=e[i].nxt) {
		if(vis[e[i].v]) continue;
		res=S-sum[e[i].v];
		ans-=sum[e[i].v],tax[col[x]]-=sum[e[i].v];
		for(re int j=0;j<v[e[i].v].size();j++) 
			ans-=v[e[i].v][j],tax[c[e[i].v][j]]-=v[e[i].v][j];
		find(e[i].v,0);
		for(re int j=0;j<v[e[i].v].size();j++) 
			ans+=v[e[i].v][j],tax[c[e[i].v][j]]+=v[e[i].v][j];
		ans+=sum[e[i].v],tax[col[x]]+=sum[e[i].v];
	}
	for(re int i=head[x];i;i=e[i].nxt) {
		if(vis[e[i].v]) continue;
		for(re int j=0;j<v[e[i].v].size();j++)
			tax[c[e[i].v][j]]-=v[e[i].v][j];
		v[e[i].v].clear(),c[e[i].v].clear();
	}
	tax[col[x]]=0;
	for(re int i=head[x];i;i=e[i].nxt) {
		if(vis[e[i].v]) continue;
		now=inf,S=sum[e[i].v],getroot(e[i].v,0),dfs(rt);
	}
}
int main() {
	n=read();int x,y;
	for(re int i=1;i<=n;i++) col[i]=read();
	for(re int i=1;i<n;i++) x=read(),y=read(),add(x,y),add(y,x);
	S=n,now=inf,getroot(1,0);dfs(rt);
	for(re int i=1;i<=n;i++) printf("%lld\n",Ans[i]+1ll);
	return 0;
}