Problem Description

Given a tree whose nodes have colors, add edges between every pair of nodes whose distance on the tree is $2$ (while keeping the original edges). Find the number of non-empty connected vertex sets in the new graph such that all vertices in the set have the same color. Since the answer may be very large, output it modulo $10^9+7$.

Definition of a connected vertex set: For a graph $G(V,E)$, a subset $V' \subseteq V$ is a connected vertex set if and only if $G(V',E')$ is a connected graph, where $E'=\{(u,v)\mid (u,v)\in E \land u \in V' \land v\in V'\}$.

Input Format

The first line contains a positive integer $n$, the number of nodes.

The next line contains $n$ positive integers. The $i$-th integer $c_i$ is the color of node $i$.

The next $n-1$ lines each contain two positive integers $u,v$, indicating that the original tree has an edge between $u$ and $v$.

Output Format

Output one integer, the answer modulo $10^9+7$.

4
1 2 1 1
1 2
2 3
2 4

8

6
1 2 2 2 1 2
5 3
2 1
4 5
6 3
3 1

14

16
1 1 2 1 1 2 2 2 1 1 2 1 1 1 2 1
12 8
14 9
10 8
1 16
7 12
6 1
14 8
3 1
12 5
1 13
12 2
1 12
15 8
11 5
4 12

442

16
11 11 11 11 11 11 11 11 11 11 11 11 11 11 11 11
4 14
4 15
12 13
2 5
7 15
10 2
15 8
15 13
9 11
13 11
3 15
8 16
6 13
1 4
10 4

27454

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

16

Hint

In Sample 1, the original tree is shown below:

After adding edges between nodes at distance $2$ on the tree, the new graph is shown below:

Then the $8$ non-empty connected vertex sets with the same color are: $\{1\},\{2\},\{3\},\{4\},\{1,3\},\{1,4\},\{3,4\},\{1,3,4\}$.

This problem uses bundled testdata.

• Special Property A: All nodes have different colors.
• Special Property B: The given tree is a star (ju-hua tree, 菊花). Specifically, the $i$-th edge connects node $1$ and node $i+1$.
• Special Property C: The given tree is a chain. Specifically, the $i$-th edge connects node $i$ and node $i+1$.
• Special Property D: All nodes have the same color.

Constraints: For all testdata, $2\leq n\leq 10^5$, $1\leq c_i\leq n$.

Translated by ChatGPT 5