题目传送门：https://www.hackerrank.com/challenges/unique-colors
感谢hzq大神找来的这道题。
考虑点分治（毕竟是路经统计），对于每一个颜色，它的贡献是独立的。我们可以在一次点分治中合在一块处理（为什么时间复杂度是对的呢，因为我们每次改动只会根据当前点的颜色进行变动，而不是所有颜色对着它都来一遍）。每次先对重心单独计算答案贡献，此时也将当前区域的各个答案贡献计算出来，并以此为基础（之后称之为基准贡献，即代码中的tot）。对于每一棵子树，我们先dfs取消掉这一片区域内贡献（为什么要取消，因为这里的贡献要留到它们那一层去解决，不然会重复计算），然后单独对这里再深搜一遍 ，具体内容如下：
出现一种颜色，如果深搜过程中这是第一次出现，我们要把目前的基准贡献扣除掉该颜色的贡献（由于进入这种深搜之前，已经取消掉这一棵子树的贡献，所以实际扣除掉了外面子树的该颜色贡献）再加上外面子树大小，从而做到维护基准贡献，然后直接加给这个点就形成当前贡献了。当然如果不是第一次出现就不用管了，因为深搜来的路上已经处理过了，然后直接加上。
这是点分治做法，时间复杂度（nlogn）。还有线性做法，即用dfs序，然后对于一个区间打上差分，最后扫一遍。
点分治代码：
1 #include<bits/stdc++.h> 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(tmp<weights) weight=x,weights=tmp; 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<n;i++) insert(); 64 sum=n; work(1); 65 for(int i=1;i<=n;i++) printf("%lld ",ans[i]); 66 return 0; 67 }