P9593 "Daily OI Round 1" Block

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'\}$.

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 one integer, the answer modulo $10^9+7$.

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.** |$\text{Subtask}$|Score|$n \le$| Special Property | Dependency | | :-----------: | :-------------:|:-----------: |:-----------: |:-----------: | |$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 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