感谢所有AC

传送门

思路

       画图得到四种情况，无色，单色，多色同链（两端点），多色不同链（多端点）。

代码

#include<iostream>
#define maxn 1000100
using namespace std;
typedef long long ll;
int Map_Color[maxn], Color_Num[maxn], Subtree_Size[maxn], Cnt_Color_Num[maxn],
EndPoint_Num[maxn], EndPoint_Size[maxn], Lastpos_Color[maxn], cnt, Head[maxn], n;
ll Ans1[maxn], Ans2[maxn];
struct Edge {
	int to, nxt;
}G[maxn << 1];
//需要注意到，链式前向星存树时需要开结点数两倍空间 
//而在存图时要开边数两倍空间
void Add_Edge(int u, int v)
{
	G[++cnt].to = v;
	G[cnt].nxt = Head[u];
	Head[u] = cnt;
}
void Dfs(int u, int fa)
{
	int Subtree_Same_Color = 0, Pos_Same_Color = 0;
	//分别用来记录以u为根时，其出现同色点的子树个数和位置
	int Visit_Color_Num = Cnt_Color_Num[Map_Color[u]];
	//记录当前访问的结点的颜色已被访问的次数
	Subtree_Size[u] = 1;
	for (int i = Head[u]; i; i = G[i].nxt)
	{
		int v = G[i].to;
		if (v == fa)
			continue;
		int Now_Color_Num = Cnt_Color_Num[Map_Color[u]];
		//遍历当前子树前该颜色出现次数
		Dfs(v, u);
		Ans1[u] += 1ll * Subtree_Size[u] * Subtree_Size[v];
		//乘法原理求单色点的方案数
		Subtree_Size[u] += Subtree_Size[v];
		//计算完方案数后再算子树大小
		if (Now_Color_Num != Cnt_Color_Num[Map_Color[u]])
			Subtree_Same_Color++, Pos_Same_Color = v;
		//如果在搜索完子树v后发现有同色点，增加子树个数并记录位置
	}
	Ans1[u] += 1ll * (n - Subtree_Size[u]) * Subtree_Size[u];
	//u结点的另一部分和u的子树部分的方案计算
	if (Visit_Color_Num || (Cnt_Color_Num[Map_Color[u]] + 1) ^ Color_Num[Map_Color[u]])
		Subtree_Same_Color++;
	//如果最初访问时该颜色已经被访问过子树个数增加（因为说明以父结点那一部分的子树中有同色点）
	//如果该结点不是最后一个该颜色结点，那么说明与u的祖先结点的子树中也有同色点
	Cnt_Color_Num[Map_Color[u]]++;
	//记录当前颜色结点出现次数
	if (Subtree_Same_Color == 1)//只有一棵子树有同色点才可能成为端点
	{
		int EndPoint_Count_Size = 0;
		//u点作为端点时该端的计算值
		if (Pos_Same_Color)
			EndPoint_Count_Size = n - Subtree_Size[Pos_Same_Color];
		else
			EndPoint_Count_Size = Subtree_Size[u];
		if (!EndPoint_Num[Map_Color[u]])
			EndPoint_Size[Map_Color[u]] = EndPoint_Count_Size;
		//如果这个端点为该颜色第一个端点，则将计算值存入待计算数组中
		else if (EndPoint_Num[Map_Color[u]] == 1)
			Ans2[Map_Color[u]] = 1ll * EndPoint_Count_Size * EndPoint_Size[Map_Color[u]];
		//如果这个端点为该颜色第二个端点，计算出答案
		EndPoint_Num[Map_Color[u]]++;
		//统计端点数量
	}
	return;
}
int main(void)
{
	ios::sync_with_stdio(0);
	cin.tie(0), cout.tie(0);
	cin >> n;
	for (int i = 1; i <= n; i++)
	{
		cin >> Map_Color[i];
		Color_Num[Map_Color[i]]++;
		Lastpos_Color[Map_Color[i]] = i;
		//分别记录颜色的映射，某种颜色总数，最后一次出现该颜色的位置（为求单颜色答案）
	}
	for (int i = 1, u, v; i < n; i++)
	{
		cin >> u >> v;
		Add_Edge(u, v);
		Add_Edge(v, u);
	}
	Dfs(1, 0);
	for (int i = 1; i <= n; i++)
	{
		if (!Color_Num[i])//没有该颜色结点
			cout << 1ll * n * (n - 1) / 2 << '\n';
		else if (Color_Num[i] == 1)//单一颜色结点
			cout << Ans1[Lastpos_Color[i]] << '\n';
		else if (EndPoint_Num[i] == 2)//有且只有两个端点
			cout << Ans2[i] << '\n';
		else//三个及以上个端点
			cout << 0 << '\n';
	}
	return 0;
}