考虑点分治（毕竟是路经统计），对于每一个颜色，它的贡献是独立的。我们可以在一次点分治中合在一块处理（为什么时间复杂度是对的呢，因为我们每次改动只会根据当前点的颜色进行变动，而不是所有颜色对着它都来一遍）。每次先对重心单独计算答案贡献，此时也将当前区域的各个答案贡献计算出来，并以此为基础（之后称之为基准贡献，即代码中的tot）。对于每一棵子树，我们先dfs取消掉这一片区域内贡献（为什么要取消，因为这里的贡献要留到它们那一层去解决，不然会重复计算），然后单独对这里再深搜一遍 ，具体内容如下：
出现一种颜色，如果深搜过程中这是第一次出现，我们要把目前的基准贡献扣除掉该颜色的贡献（由于进入这种深搜之前，已经取消掉这一棵子树的贡献，所以实际扣除掉了外面子树的该颜色贡献）再加上外面子树大小，从而做到维护基准贡献，然后直接加给这个点就形成当前贡献了。当然如果不是第一次出现就不用管了，因为深搜来的路上已经处理过了，然后直接加上。
1 #include
2 using namespace std;
3 #define N 100005
4 #define INF 1e9
5 #define LL long long
6 inline int read(){
7 int x=0,f=1; char a=getchar();
8 while(a<'0' || a>'9') {if(a=='-') f=-1; a=getchar();}
9 while(a>='0' && a<='9') x=x*10+a-'0',a=getchar();
10 return x*f;
11 }
12 int n,cnt,sum,size[N],head[N],a[N],s[N],weight,weights,tot;
13 LL ans[N];
14 bool vis[N],app[N];
15 struct edges{
16 int to,next;
17 }e[2*N];
18 inline void insert(){
19 int u=read(),v=read();
20 e[cnt]=(edges){v,head[u]};head[u]=cnt++;
21 e[cnt]=(edges){u,head[v]};head[v]=cnt++;
22 }
23 void getroot(int x,int fa){
24 size[x]=1; int tmp=0;
25 for(int i=head[x];i>=0;i=e[i].next){
26 if(vis[e[i].to] || fa==e[i].to) continue;
27 getroot(e[i].to,x); size[x]+=size[e[i].to];
28 tmp=max(tmp,size[e[i].to]);
29 }
30 tmp=max(tmp,sum-size[x]);
31 if(tmptmp;
32 }
33 void dfs(int x,int fa,int f){
34 bool ok=0;
35 if(!app[a[x]] && a[x]!=a[weight]) app[a[x]]=ok=1,s[a[x]]+=size[x]*f,tot+=size[x]*f;
36 for(int i=head[x];i>=0;i=e[i].next) if(!vis[e[i].to] && fa!=e[i].to) dfs(e[i].to,x,f);
37 if(ok) app[a[x]]=0;
38 }
39 void DFS(int x,int fa,int p){
40 bool ok=0;
41 if(!app[a[x]] && a[x]!=a[weight]) app[a[x]]=ok=1,tot+=p-s[a[x]];
42 ans[x]+=tot;
43 for(int i=head[x];i>=0;i=e[i].next) if(!vis[e[i].to] && fa!=e[i].to) DFS(e[i].to,x,p);
44 if(ok) app[a[x]]=0,tot-=p-s[a[x]];
45 }
46 void work(int x){
47 weights=INF; getroot(x,0); getroot(weight,0);
48 x=weight; vis[x]=1;
49 dfs(weight,0,1); tot+=size[x]; ans[x]+=tot;
50 for(int i=head[x];i>=0;i=e[i].next){
51 if(vis[e[i].to]) continue;
52 dfs(e[i].to,x,-1); tot-=size[e[i].to];
53 DFS(e[i].to,x,size[x]-size[e[i].to]);
54 dfs(e[i].to,x,1); tot+=size[e[i].to];
55 }
56 dfs(weight,0,-1); tot=0;
57 for(int i=head[x];i>=0;i=e[i].next)
58 if(!vis[e[i].to]) sum=size[e[i].to],work(e[i].to);
59 }
60 int main(){
61 n=read(); memset(head,-1,sizeof(head));
62 for(int i=1;i<=n;i++) a[i]=read();
63 for(int i=1;i) insert();
64 sum=n; work(1);
65 for(int i=1;i<=n;i++) printf("%lld\n",ans[i]);
66 return 0;
67 }
京公网安备 11010802041100号