lrb有一棵树,树的每个节点有个颜色。给一个长度为n的颜色序列,定义s(i,j)s(i,j)iijj 的颜色数量。以及

f(i)=i=1ns(i,j)f(i) = \sum_{i=1}^n s(i,j)

求出所有f(i)f(i)

链接🔗

「P2664」 树上游戏

题解

我们可以用某种颜色的点将树切割成块,块内任意两点不会有该颜色的贡献
所以我们可以得到以下结论
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;
}