题目描述
Problem Description
Given a tree where each node has a color, edges are added between nodes that are originally two steps apart (while retaining the original edges). The task is to find the number of non-empty sets of nodes that are both the same color and connected in the new graph. Since the answer may be very large, you only need to output the answer modulo $10^9+7$.
Definition of a connected set of nodes: For a graph $G(V,E)$, a subset $V'$ of $V$ is a connected set of nodes if and only if $G(V',E')$ is a connected graph, where the edge set $E'=\{(u,v)|(u,v)\in E \land u \in V' \land v \in V'\}$.
Input Format
The first line contains a positive integer $n$, representing the number of nodes.
The next line contains $n$ positive integers, where the $i$-th integer $c_i$ represents the color of the $i$-th node.
The following $n-1$ lines each contain two positive integers $u$ and $v$, indicating that there is an edge in the original tree from node $u$ to node $v$.
Output Format
A single line containing an integer, representing the answer modulo $10^9+7$.
Sample Input and Output
Sample Input #1
4
1 2 1 1
1 2
2 3
2 4
Sample Output #1
8
Sample Input #2
6
1 2 2 2 1 2
5 3
2 1
4 5
6 3
3 1
Sample Output #2
14
Sample Input #3
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
Sample Output #3
442
Sample Input #4
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
Sample Output #4
27454
Sample Input #5
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
Sample Output #5
16
Notes
In Sample 1, the original tree is shown below:
After adding edges between nodes that are two steps apart, the new graph is shown below:
The 8 sets of nodes that are the same color and connected are: $\{1\},\{2\},\{3\},\{4\},\{1,3\},\{1,4\},\{3,4\},\{1,3,4\}$.
This problem has bundled test cases.
| Subtask | Points | $n \le$ | Special Conditions | Subtask Dependencies |
|---|---|---|---|---|
| $0$ | $5$ | $10^5$ | A | None |
| $1$ | $5$ | $16$ | None | None |
| $2$ | $5$ | $10^5$ | B | None |
| $3$ | $15$ | $10^5$ | C | None |
| $4$ | $20$ | $10^5$ | D | None |
| $5$ | $50$ | $10^5$ | None | $0\sim4$ |
- Special Condition A: All nodes have different colors.
- Special Condition B: The given tree is a star, specifically, the $i$-th edge connects node $1$ and node $i+1$.
- Special Condition C: The given tree is a chain, specifically, the $i$-th edge connects node $i$ and node $i+1$.
- Special Condition D: All nodes have the same color.
For all data, $2 \leq n \leq 10^5$, $1 \leq c_i \leq n$.
输入格式
见题面 Input 部分
输出格式
见题面 Output 部分