AT_abc359_g [ABC359G] Sum of Tree Distance

题目描述

给定一棵有 $N$ 个顶点的树。第 $i$ 条边连接顶点 $u_i$ 和顶点 $v_i$，是双向的。

另外，给定一个整数序列 $A=(A_1,\ldots,A_N)$。

这里定义 $f(i,j)$ 如下：

• 当 $A_i=A_j$ 时，$f(i,j)$ 为从顶点 $i$ 到顶点 $j$ 需要经过的最少边数。
• 当 $A_i\neq A_j$ 时，$f(i,j)=0$。

请计算下式的值：

输入格式

输入按以下格式从标准输入读入。

$N$
$u_1$ $v_1$
$\vdots$
$u_{N-1}$ $v_{N-1}$
$A_1$ $A_2$ $\ldots$ $A_N$

输出格式

输出答案。

输入输出样例 #1

输入 #1

4
3 4
4 2
1 2
2 1 1 2

输出 #1

4

输入输出样例 #2

输入 #2

8
8 6
3 8
1 4
7 8
4 5
3 4
8 2
1 2 2 2 3 1 1 3

输出 #2

19

说明/提示

限制条件

• $2\leq N\leq 2\times 10^5$
• $1\leq u_i,v_i \leq N$
• $1\leq A_i\leq N$
• 输入的图为一棵树
• 输入的所有数均为整数

样例解释 1

有 $f(1,4)=2, f(2,3)=2$。除此之外，满足 $1\leq i<j\leq N$ 的其它 $i,j$，$f(i,j)=0$，所以答案为 $2+2=4$。

由 ChatGPT 4.1 翻译