题目传送门: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 ;
}