洛谷 P2664 树上游戏 题解

题目传送门

题目大意： 求出所有的 $sum_i=\sum_{j=1}^n s(i,j)$，其中 $s(i,j)$ 表示 $i$ 到 $j$ 的路径上有多少种不同的颜色。

题解

树上的路径的统计类型的题目肯定是用点分治来搞嘛。对于当前分治到的一颗子树，我们考虑统计所有子树内经过重心的路径所产生的贡献。

为了方便，先定义几个东西：

1. $col[i]$ 表示节点 $i$ 的颜色
2. $val[i]$ 表示 $i$ 这种颜色给重心带来的贡献
3. $size[i]$ 表示以 $i$ 为根的子树的大小
4. $root$ 表示重心

那么我们在从重心遍历下去的过程中，对于遍历到的一个点 $x$，假如这个点的颜色在它到重心的路径上是第一次出现，那么我们让 $val[col[x]]$ 加上 $size[x]$，这个很容易想明白，就不多解释了。

遍历完一次后， $tot=\sum_i val[i]$ 就是重心得到的贡献。然后我们利用这个 $tot$ 和 $val$ 数组来统计其它节点得到的贡献。

对于重心的一棵以 $y$ 为根的子树，统计贡献时，我们首先将这颗子树对 $val$ 和 $tot$ 的贡献消去，然后对子树内每个结点分类讨论：

1. 假如颜色 $c$ 在它到重心的路径上，那么颜色 $c$ 给他产生的贡献就是 $size[root]-size[y]$
2. 假如颜色 $c$ 不在它到重心的路径上，那么颜色 $c$ 给它产生的贡献就是 $val[c]$。

将上面两个情况综合一下，统一利用 $tot$ 来统计贡献——我们在从 $root$ 遍历下来时，假如遇到了一种第一次出现的颜色 $c$，那么我们让 $tot$ 减去 $val[c]$ 然后加上 $size[root]-size[y]$。这样的话，在遍历到每个点的时候，此时的 $tot$ 就是它获得的贡献。

统计完一棵子树后，记得还原它对 $tot$ 和 $val$ 的贡献。

代码如下：

#include <cstdio>
#define maxn 100010
#define inf 999999999
#define ll long long 

int n,col[maxn];
struct edge{int y,next;};
edge e[maxn<<1];
int first[maxn],len=0;
void buildroad(int x,int y)
{
	e[++len]=(edge){y,first[x]};
	first[x]=len;
}
bool v[maxn];
int Size,root,rt_mson,size[maxn],mson[maxn];
void getroot(int x,int fa)//找重心
{
	size[x]=1;mson[x]=0;
	for(int i=first[x];i;i=e[i].next)
	{
		int y=e[i].y;
		if(v[y]||y==fa)continue;
		getroot(y,x);
		size[x]+=size[y];
		if(size[y]>mson[x])mson[x]=size[y];
	}
	if(Size-size[x]>mson[x])mson[x]=Size-size[x];
	if(mson[x]<rt_mson)rt_mson=mson[x],root=x;
}
ll sum[maxn],val[maxn],tot;
bool color[maxn];//color[i] 表示颜色i是否出现过
void get_size(int x,int fa)//求size数组
{
	size[x]=1;val[col[x]]=0;color[col[x]]=false;
	for(int i=first[x];i;i=e[i].next)
	{
		int y=e[i].y;
		if(v[y]||y==fa)continue;
		get_size(y,x);
		size[x]+=size[y];
	}
}
void solve_1(int x,int fa,int type)//统计 tot 和 val，type表示是删去贡献还是加上贡献
{
	bool tf=false;//回溯时用到
	if(!color[col[x]])color[col[x]]=tf=true,tot+=size[x]*type,val[col[x]]+=size[x]*type;
	for(int i=first[x];i;i=e[i].next)
	{
		int y=e[i].y;
		if(v[y]||y==fa)continue;
		solve_1(y,x,type);
	}
	if(tf)color[col[x]]=false;
}
void solve_2(int x,int fa,int Siz)//统计sum
{
	bool tf=false;
	//假如颜色col[x]是第一次出现，那么让更新tot
	if(!color[col[x]])color[col[x]]=tf=true,tot-=val[col[x]]-Siz;
	sum[x]+=tot;//累计贡献
	for(int i=first[x];i;i=e[i].next)
	{
		int y=e[i].y;
		if(v[y]||y==fa)continue;
		solve_2(y,x,Siz);
	}
	if(tf)color[col[x]]=false,tot+=val[col[x]]-Siz;
}
void conquer(int rt)
{
	get_size(rt,0);
	tot=0;solve_1(rt,0,1);
	
	sum[rt]+=tot;
	color[col[rt]]=true;tot-=val[col[rt]];
	for(int i=first[rt];i;i=e[i].next)
	if(!v[e[i].y])
	{
		tot+=size[rt]-size[e[i].y];//这一步和上面的 tot-=val[col[rt]] 都是用来统计重心的贡献的
		solve_1(e[i].y,rt,-1);//先消去这棵子树的贡献
		solve_2(e[i].y,rt,size[rt]-size[e[i].y]);
		solve_1(e[i].y,rt,1);//还原贡献
		tot-=size[rt]-size[e[i].y];
	}
}
void divide(int rt)//点分治
{
	v[rt]=true;conquer(rt);
	for(int i=first[rt];i;i=e[i].next)
	{
		int y=e[i].y;
		if(v[y])continue;
		Size=size[y];root=0;rt_mson=inf;
		getroot(y,rt); divide(root);
	}
}

int main()
{
	scanf("%d",&n);
	for(int i=1;i<=n;i++)
	scanf("%d",&col[i]);
	for(int i=1,x,y;i<n;i++)
	scanf("%d %d",&x,&y),buildroad(x,y),buildroad(y,x);
	Size=n;root=0;rt_mson=inf;
	getroot(1,0);
	divide(root);
	for(int i=1;i<=n;i++)
	printf("%lld\n",sum[i]);
}