lrb有一棵树,树的每个节点有个颜色。给一个长度为n的颜色序列,定义 为 到 的颜色数量。以及
求出所有
链接🔗
题解
我们可以用某种颜色的点将树切割成块,块内任意两点不会有该颜色的贡献
所以我们可以得到以下结论
1. 某个颜色对于一个点的贡献值为 n - size
2. 某一个点,所有颜色对其的贡献为,n * 颜色数 - 每种颜色节点时所在块的节点数
因为每个点的颜色,对于每个块,我们可以通过差分的思想,得到块的大小 (dfs) 函数
val 表示以该节点父亲为切割点时,该节点的块的大小
root 表示当前所在节点,所有包含次节点的颜色块大小
num 表示第i个颜色块,当前点的贡献
代码
#include <cstdio>
#include <cstring>
#include <algorithm>
using namespace std;
const int maxd = 1e5+10;
struct node
{
int v,net;
}e[maxd<<1];
int f[maxd],eid;
void insert(int x,int y)
{
e[eid].net = f[x];
e[eid].v = y;
f[x] = eid++;
}
void init()
{
memset(f,-1,sizeof(f));
eid = 0;
}
int n,sum; //总颜色数,根结点被所有颜色块包含块的总大小
int sz[maxd],vis[maxd],col[maxd],del[maxd],val[maxd],num[maxd];
// 子树大小,该颜色是否出现过,颜色,该点颜色删除点数,所在点被父亲颜色切割的块的大小,根所在块被分割的大小
long long ans[maxd],root;
void get_ans(int x,int fa=0)
{
long long tmp = num[col[fa]];
root += 1ll * val[x] - num[col[fa]];
num[col[fa]] = val[x];
ans[x] = 1ll * n * sum - root + num[col[x]]; // 因为当前点一定会被包含,所以要加上当前颜色
for(int i = f[x];i+1;i = e[i].net)
{
int v = e[i].v;
if(v == fa) continue;
get_ans(v,x);
}
num[col[fa]] = tmp;
root -= val[x] - num[col[fa]];
}
void dfs(int x,int fa = 0)
{
sz[x] = 1;
int tmp = del[col[fa]]; // 在这之前就已经删除的 父亲颜色
for(int i=f[x];i+1;i=e[i].net)
{
int v = e[i].v;
if(v == fa) continue;
dfs(v,x);
sz[x] += sz[v];
}
del[col[x]]++; // 该点颜色删除了
if(fa)
{
val[x] = sz[x] - del[col[fa]] + tmp; // 减去所以删除的点 + 之前删除的点,就是该子树下删除的点
del[col[fa]] += val[x]; // 将该区块的点加入删除的点
}
}
int main()
{
// freopen("a.in","r",stdin);
// freopen("k.out","w",stdout);
init();
scanf("%d",&n);
for(int i=1;i<=n;i++)
{
scanf("%d",&col[i]);
if(!vis[col[i]]) vis[col[i]]=1,sum++;
}
for(int i=1;i<n;i++)
{
int x,y; scanf("%d %d",&x,&y);
insert(x,y); insert(y,x);
}
dfs(1);
for(int i=1;i<maxd;i++)
if(vis[i]) root += 1ll * n - del[i],num[i] = n-del[i]; //
get_ans(1);
for(int i=1;i<=n;i++) printf("%lld\n",ans[i]);
return 0;
}