题解 P5593 【小猪佩奇爬树 加强版】

题解 - $\mathrm{P}5593$

题目描述

题目传送门

给出一棵 $n$ 个点的树，每个点上有一种颜色。现在请你求出对于每一种颜色，树上有多少条链包含该种颜色的所有点。

$n\le 3\times 10^6$

$\mathrm{S}\mathrm{o}\mathrm{l}$

一道细节很多的题目。

首先我们很容易想到若一种颜色的数量 $t_i=0$ 那么输出 $\frac{n\times (n+1)}{2}$。已经如果一种颜色在大于等于三棵子树中出现过那么输出 $0$。

于是我们考虑 $t_i=1$ 的情况，这个也比较容易答案就是这个节点为根的子树大小相互乘起来再累加。

难点在于某种颜色有两个端点的情况的处理。我们记 $m{p}_c$ 为颜色 $c$ 到现在出现的个数，并且记录 $tot$ 表示在当前子树中有几个颜色为 $c$ 的个数。那么每次递归完一颗子树看 $m{p}_c$ 是否改变，若改变那么 $tot$ 加一。若 $tot=1$ 则表示其为该颜色的某一个端点。那么答案即为两个端的子树大小乘积，具体实现看代码。

时间复杂度 $O(n)$

$\mathrm{C}\mathrm{o}\mathrm{d}\mathrm{e}$

const int N=3e6+5;

int n,m,cnt,t[N],col[N],siz[N],Siz[N];
int gs[N],cs[N],head[N],Nod[N],Ans[N];

struct Node
{
    
    
	int nex,to;
};
Node e[N<<1];

inline void jia(int u,int v)
{
    
    
	e[++cnt].nex=head[u];
	head[u]=cnt;
	e[cnt].to=v;
}

inline void dfs(int u,int fa)
{
    
    
	siz[u]=1;
	int Col=col[u],sum=0;
	int onend=0,Las=cs[Col];
	GO(i,u) 
	{
    
    
		int v=e[i].to;
		if(v==fa) continue;
		int las=cs[Col];//递归这颗子树之前的颜色个数
		dfs(v,u);
		Siz[u]+=siz[u]*siz[v];//对就1个该颜色答案的统计
		siz[u]+=siz[v];
		if(las^cs[Col]) sum++,onend=v;//记录端点
	}
	Siz[u]+=siz[u]*(n-siz[u]);
	if(Las||t[Col]-1!=cs[Col]) sum++;//子树u自己本身也要统计进去
	cs[Col]++;
	if(sum==1) 
	{
    
    
		if(!gs[Col]) Nod[Col]=u;
		else 
		{
    
    
			int mul=(onend)?n-siz[onend]:siz[u];
			Ans[Col]=siz[Nod[Col]]*mul;
		}
		gs[Col]++;//统计有几个子树有颜色Col
	}
}

signed main()
{
    
    
	io.read(n);
	For(i,1,n) 	
	{
    
    
		io.read(col[i]);
		t[col[i]]++;
		Nod[col[i]]=i;
	}
	For(i,1,n-1)
	{
    
    
		int x,y;
		io.read(x),io.read(y);
		jia(x,y),jia(y,x);
	}
	dfs(1,0);
	For(i,1,n) 
	{
    
    
		int ret=0;
		if(!t[i]) ret=n*(n-1)/2ll;//没有这种颜色
		if(t[i]==1) ret=Siz[Nod[i]]; //有1种
		if(gs[i]==2) ret=Ans[i];//两个端点情况
		io.write(ret),puts("");
	}
	return 0;
}