日常小测：颜色 && Hackerrank Unique_colors

题目传送门：https://www.hackerrank.com/challenges/unique-colors

感谢hzq大神找来的这道题。

考虑点分治（毕竟是路经统计），对于每一个颜色，它的贡献是独立的。我们可以在一次点分治中合在一块处理（为什么时间复杂度是对的呢，因为我们每次改动只会根据当前点的颜色进行变动，而不是所有颜色对着它都来一遍）。每次先对重心单独计算答案贡献，此时也将当前区域的各个答案贡献计算出来，并以此为基础（之后称之为基准贡献，即代码中的tot）。对于每一棵子树，我们先dfs取消掉这一片区域内贡献（为什么要取消，因为这里的贡献要留到它们那一层去解决，不然会重复计算），然后单独对这里再深搜一遍 ，具体内容如下：

出现一种颜色，如果深搜过程中这是第一次出现，我们要把目前的基准贡献扣除掉该颜色的贡献（由于进入这种深搜之前，已经取消掉这一棵子树的贡献，所以实际扣除掉了外面子树的该颜色贡献）再加上外面子树大小，从而做到维护基准贡献，然后直接加给这个点就形成当前贡献了。当然如果不是第一次出现就不用管了，因为深搜来的路上已经处理过了，然后直接加上。

这是点分治做法，时间复杂度（nlogn）。还有线性做法，即用dfs序，然后对于一个区间打上差分，最后扫一遍。

点分治代码：

 #include<bits/stdc++.h>
 using namespace std;
 #define N 100005
 #define INF 1e9
 #define LL long long
 inline int read(){
     int x=,f=; char a=getchar();
     while(a<'' || a>'') {if(a=='-') f=-; a=getchar();}
     while(a>='' && a<='') x=x*+a-'',a=getchar();
     return x*f;
 }
 int n,cnt,sum,size[N],head[N],a[N],s[N],weight,weights,tot;
 LL ans[N];
 bool vis[N],app[N];
 struct edges{
     int to,next;
 }e[*N];
 inline void insert(){
     int u=read(),v=read();
     e[cnt]=(edges){v,head[u]};head[u]=cnt++;
     e[cnt]=(edges){u,head[v]};head[v]=cnt++;
 }
 void getroot(int x,int fa){
     size[x]=; int tmp=;
     for(int i=head[x];i>=;i=e[i].next){
         if(vis[e[i].to] || fa==e[i].to) continue;
         getroot(e[i].to,x); size[x]+=size[e[i].to];
         tmp=max(tmp,size[e[i].to]);
     }
     tmp=max(tmp,sum-size[x]);
     if(tmp<weights) weight=x,weights=tmp;
 }
 void dfs(int x,int fa,int f){
     bool ok=;
     if(!app[a[x]] && a[x]!=a[weight]) app[a[x]]=ok=,s[a[x]]+=size[x]*f,tot+=size[x]*f;
     for(int i=head[x];i>=;i=e[i].next) if(!vis[e[i].to] && fa!=e[i].to) dfs(e[i].to,x,f);
     if(ok) app[a[x]]=;
 }
 void DFS(int x,int fa,int p){
     bool ok=;
     if(!app[a[x]] && a[x]!=a[weight]) app[a[x]]=ok=,tot+=p-s[a[x]];
     ans[x]+=tot;
     for(int i=head[x];i>=;i=e[i].next) if(!vis[e[i].to] && fa!=e[i].to) DFS(e[i].to,x,p);
     if(ok) app[a[x]]=,tot-=p-s[a[x]];
 }
 void work(int x){
     weights=INF; getroot(x,); getroot(weight,);
     x=weight; vis[x]=;
     dfs(weight,,); tot+=size[x]; ans[x]+=tot;
     for(int i=head[x];i>=;i=e[i].next){
         if(vis[e[i].to]) continue;
         dfs(e[i].to,x,-); tot-=size[e[i].to];
         DFS(e[i].to,x,size[x]-size[e[i].to]);
         dfs(e[i].to,x,); tot+=size[e[i].to];
     }
     dfs(weight,,-); tot=;
     for(int i=head[x];i>=;i=e[i].next)
     if(!vis[e[i].to]) sum=size[e[i].to],work(e[i].to);
 }
 int main(){
     n=read(); memset(head,-,sizeof(head));
     for(int i=;i<=n;i++) a[i]=read();
     for(int i=;i<n;i++) insert();
     sum=n; work();
     for(int i=;i<=n;i++) printf("%lld\n",ans[i]);
     return ;
 }