G - Sum of Tree Distance Editorial /

Time Limit: 4 sec / Memory Limit: 1024 MiB

配点 : 600 点

問題文

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 とする。

次の式の値を求めてください。

\displaystyle \sum_{i=1}^{N-1}\sum_{j=i+1}^N f(i,j)


制約

  • 2\leq N\leq 2\times 10^5
  • 1\leq u_i,v_i \leq N
  • 1\leq A_i\leq N
  • 入力されるグラフは木
  • 入力される数値は全て整数

入力

入力は以下の形式で標準入力から与えられる。

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

出力

答えを出力せよ。

入力例 1

4
3 4
4 2
1 2
2 1 1 2

出力例 1

4

f(1,4)=2,f(2,3)=2 となります。また、それ以外の i,j\ (1\leq i < j\leq N) について f(i,j)=0 なので、答えは 2+2=4 です。

入力例 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

Score : 600 points

Problem Statement

You are given a tree with N vertices. The i-th edge connects vertices u_i and v_i bidirectionally.

Additionally, you are given an integer sequence A=(A_1,\ldots,A_N).

Here, define f(i,j) as follows:

  • If A_i = A_j, then f(i,j) is the minimum number of edges you need to traverse to move from vertex i to vertex j. If A_i \neq A_j, then f(i,j) = 0.

Calculate the value of the following expression:

\displaystyle \sum_{i=1}^{N-1}\sum_{j=i+1}^N f(i,j).


Constraints

  • 2 \leq N \leq 2 \times 10^5
  • 1 \leq u_i, v_i \leq N
  • 1 \leq A_i \leq N
  • The input graph is a tree.
  • All input values are integers.

Input

The input is given from Standard Input in the following format:

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

Output

Print the answer.

Sample Input 1

4
3 4
4 2
1 2
2 1 1 2

Sample Output 1

4

f(1,4)=2, f(2,3)=2. For all other i, j\ (1 \leq i < j \leq N), we have f(i,j)=0, so the answer is 2+2=4.

Sample Input 2

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

Sample Output 2

19