题面:
题解:
点分治。
考虑当前点分治子树的根节点为p。
那么我们需要处理两部分贡献。
①、所有节点对点p的贡献(包括p节点自己)
②、经过点p的路径(x,y)对x点的贡献。
第一类贡献:
如果一个节点 x 的颜色 color [ x ] 在它到根节点 p 的路径上是第一次出现,那么以 x 为根的子树上的每个节点,都可以对节点 p 产生贡献,那么一***生了 size [ x ] 的贡献。
为方便我们最后计算和统计答案,我们把所有这样的贡献加起来记为sum。显然该当前分治子树上所有节点对p的贡献和就是sum。即 ans[p]+=sum。
void dfs_sum(int x,int fa)
{
si[x]=1;
cnt[vi[x]]++;
for(int i=head[x];i;i=nt[i])
{
int y=ver[i];
if(ha[y]||y==fa) continue;
dfs_sum(y,x);
si[x]+=si[y];
}
if(cnt[vi[x]]==1)
{
sum+=si[x];
c[vi[x]]+=si[x];
}
cnt[vi[x]]--;
}
第二类贡献。
假设当前考虑点经过p的路径对点x的贡献
我们将第二类贡献分为两部分
第一部分:
路径p–x出现的颜色对x点的贡献(不包括p点)。这一部分的贡献为 (size[ p ] - size[ sonp ] ) * res
其中 sonp 为 p 的儿子且为 x 的祖先,res为p–x出现的颜色种类数(不包括p点)。
第二部分:
除了 sonp 以外子树上新出现的颜色对 x 的贡献,我们求出了sum,所有节点对p的贡献,那么我们需要减去sonp子树中对sum的贡献,然后再减去p–x出现的颜色对sum的贡献(不包括p点),此时的sum就是sonp以外子树新出现的颜色对x的贡献
void dfs_ans(int x,int fa)
{
cnt[vi[x]]++;
if(cnt[vi[x]]==1)
{
sum-=c[vi[x]];
res++;
}
ans[x]+=sum+res*six_xiy;
for(int i=head[x];i;i=nt[i])
{
int y=ver[i];
if(y==fa||ha[y]) continue;
dfs_ans(y,x);
}
if(cnt[vi[x]]==1)
{
sum+=c[vi[x]];
res--;
}
cnt[vi[x]]--;
}
代码:
#include<iostream>
#include<cstdio>
#include<cstring>
#include<string>
#include<cmath>
#include<algorithm>
#include<cstdlib>
#define ll long long
#define llu unsigned ll
#define pr make_pair
#define pb push_back
using namespace std;
const int maxn=100100;
const int inf=0x3f3f3f3f;
int head[maxn],ver[maxn<<1],nt[maxn<<1];
int max_part[maxn],si[maxn],vi[maxn],c[maxn],cnt[maxn];
bool ha[maxn];
int tot=1,root,max_size,n;
ll ans[maxn],sum,res,six_xiy;
void add(int x,int y)
{
ver[++tot]=y,nt[tot]=head[x],head[x]=tot;
}
void dfs_size(int x,int fa)
{
si[x]=1;max_part[x]=0;
for(int i=head[x];i;i=nt[i])
{
int y=ver[i];
if(ha[y]||y==fa) continue;
dfs_size(y,x);
si[x]+=si[y];
max_part[x]=max(max_part[x],si[y]);
}
}
void dfs_root(int now_root,int x,int fa)
{
max_part[x]=max(max_part[x],si[now_root]-si[x]);
if(max_size>max_part[x])
{
max_size=max_part[x];
root=x;
}
for(int i=head[x];i;i=nt[i])
{
int y=ver[i];
if(ha[y]||y==fa) continue;
dfs_root(now_root,y,x);
}
}
void dfs_sum(int x,int fa)
{
si[x]=1;
cnt[vi[x]]++;
for(int i=head[x];i;i=nt[i])
{
int y=ver[i];
if(ha[y]||y==fa) continue;
dfs_sum(y,x);
si[x]+=si[y];
}
if(cnt[vi[x]]==1)
{
sum+=si[x];
c[vi[x]]+=si[x];
}
cnt[vi[x]]--;
}
void change(int x,int fa,int val)
{
cnt[vi[x]]++;
for(int i=head[x];i;i=nt[i])
{
int y=ver[i];
if(y==fa||ha[y]) continue;
change(y,x,val);
}
if(cnt[vi[x]]==1)
{
sum+=si[x]*val;
c[vi[x]]+=si[x]*val;
}
cnt[vi[x]]--;
}
void dfs_ans(int x,int fa)
{
cnt[vi[x]]++;
if(cnt[vi[x]]==1)
{
sum-=c[vi[x]];
res++;
}
ans[x]+=sum+res*six_xiy;
for(int i=head[x];i;i=nt[i])
{
int y=ver[i];
if(y==fa||ha[y]) continue;
dfs_ans(y,x);
}
if(cnt[vi[x]]==1)
{
sum+=c[vi[x]];
res--;
}
cnt[vi[x]]--;
}
void dfs_clear(int x,int fa)
{
cnt[vi[x]]=0;
c[vi[x]]=0;
for(int i=head[x];i;i=nt[i])
{
int y=ver[i];
if(y==fa||ha[y]) continue;
dfs_clear(y,x);
}
}
void do_it(int x)
{
sum=0,res=0;
dfs_sum(x,0);
ans[x]+=sum;
for(int i=head[x];i;i=nt[i])
{
int y=ver[i];
if(ha[y]) continue;
cnt[vi[x]]++;
sum-=si[y];
c[vi[x]]-=si[y];
change(y,x,-1);
cnt[vi[x]]--;
six_xiy=si[x]-si[y];
dfs_ans(y,x);
cnt[vi[x]]++;
sum+=si[y];
c[vi[x]]+=si[y];
change(y,x,1);
cnt[vi[x]]--;
}
dfs_clear(x,0);
}
void dfs(int x)
{
max_size=inf;
dfs_size(x,0);
dfs_root(x,x,0);
do_it(root);
ha[root]=true;
for(int i=head[root];i;i=nt[i])
{
int y=ver[i];
if(ha[y]) continue;
dfs(y);
}
}
int main(void)
{
scanf("%d",&n);
for(int i=1;i<=n;i++)
scanf("%d",&vi[i]);
int x,y;
for(int i=1;i<n;i++)
scanf("%d%d",&x,&y),add(x,y),add(y,x);
dfs(1);
for(int i=1;i<=n;i++)
printf("%lld\n",ans[i]);
return 0;
}