这真是一道神仙的一批的题目
定义(s(i,j))表示从点(i)到点(j)经过的颜色数量
设
[sum_i=sum_{j=1}^ns(i,j)
]
求出所有的(sum_i)
考虑点分治
对于一个点我们用两种方式来统计其答案
-
这个点作为分治重心时，分值区域内所有点到这个点贡献
-
这个点不是分治重心的时候，当前分治区域内其他子树到这个点的贡献
第一种贡献我们很好统计，点分治的时候把所有子树遍历一遍就好了
第二种就需要转换一下思路了，我们不能直接求(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
",Ans[i]+1ll);
return 0;
}