正题

P2664

---

题目大意

定义s(i,j)为i到j路径上的颜色数量，$su{m}_i={\sum }_{j=1}^{n}s(i,j)$

现在让你求sum的和

---

解题思路

考虑一种颜色的贡献：

先把该颜色的点删掉，这样就形成了若干小树

不难发现，同一小树中的点对不会产生贡献，而不同小树的点对会产生一点贡献

那么在计算一个点的贡献时，就是n-所在小树的size

那么直接dfs遍历每一个点，然后计算

因为一个点只会影响当前颜色的贡献，所以时间是O(n)的

---

code

#include<cstdio>
#include<cstring>
#include<iostream>
#include<algorithm>
#define ll long long
#define N 100010
using namespace std;
ll n,x,y,now,tot,sum,f[N],p[N],h[N],c[N],sz[N],num[N],ans[N];
struct rec
{ll to,nx;
}e[N<<1];
void add(ll x,ll y)
{e[++tot].to=y;e[tot].nx=h[x];h[x]=tot;return;
}
void dfs1(ll x,ll fa)
{ll g=num[c[fa]];num[c[x]]++;sz[x]=1;for(ll i=h[x];i;i=e[i].nx){ll y=e[i].to;if(y==fa)continue;dfs1(y,x);sz[x]+=sz[y];}if(fa)num[c[fa]]+=(f[x]=sz[x]-(num[c[fa]]-g));//计算小树大小return;
}
void dfs2(ll x,ll fa)
{ll g=num[c[fa]];now+=f[x]-num[c[fa]];//减去上一小树大小num[c[fa]]=f[x];ans[x]=sum*n-now+num[c[x]];//当前结点颜色能不减for(ll i=h[x];i;i=e[i].nx){ll y=e[i].to;if(y==fa)continue;dfs2(y,x);}num[c[fa]]=g;now-=f[x]-num[c[fa]];return;
}
int main()
{scanf("%lld",&n);for(ll i=1;i<=n;++i){scanf("%lld",&c[i]);if(!p[c[i]])p[c[i]]=1,sum++;}for(ll i=1;i<n;++i){scanf("%lld%lld",&x,&y);add(x,y);add(y,x);}dfs1(1,0);num[0]=0;for(ll i=1;i<=100000;++i){if(!p[i])continue;num[i]=n-num[i];now+=num[i];}dfs2(1,0);for(ll i=1;i<=n;++i)printf("%lld\n",ans[i]);return 0;
}