日常小测：颜色 && Hackerrank Unique_colors

Posted 2020-08-25 Izaya

题目传送门：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\n",ans[i]);
66     return 0;
67 }