Count Paths
Contest: Educational Codeforces Round 162 (Rated for Div. 2)
Difficulty rating: 2000
Tags: data structures, dfs and similar, dp, dsu, graphs, trees

You are given a tree, consisting of $$$n$$$ vertices, numbered from $$$1$$$ to $$$n$$$. Every vertex is colored in some color, denoted by an integer from $$$1$$$ to $$$n$$$.

A simple path of the tree is called beautiful if:

- it consists of at least $$$2$$$ vertices;
- the first and the last vertices of the path have the same color;
- no other vertex on the path has the same color as the first vertex.

Count the number of the beautiful simple paths of the tree. Note that paths are considered undirected (i. e. the path from $$$x$$$ to $$$y$$$ is the same as the path from $$$y$$$ to $$$x$$$).

Input format:
The first line contains a single integer $$$t$$$ ($$$1 \le t \le 10^4$$$) — the number of testcases.

The first line of each testcase contains a single integer $$$n$$$ ($$$2 \le n \le 2 \cdot 10^5$$$) — the number of vertices in the tree.

The second line contains $$$n$$$ integers $$$c_1, c_2, \dots, c_n$$$ ($$$1 \le c_i \le n$$$) — the color of each vertex.

The $$$i$$$-th of the next $$$n - 1$$$ lines contains two integers $$$v_i$$$ and $$$u_i$$$ ($$$1 \le v_i, u_i \le n$$$; $$$v_i \neq u_i$$$) — the $$$i$$$-th edge of the tree.

The given edges form a valid tree. The sum of $$$n$$$ over all testcases doesn't exceed $$$2 \cdot 10^5$$$.

Output format:
For each testcase, print a single integer — the number of the beautiful simple paths of the tree.

Constraints:
Time limit: 2.0 s
Memory limit: 512.0 MB

Example input:
4
3
1 2 1
1 2
2 3
5
2 1 2 1 2
1 2
1 3
3 4
4 5
5
1 2 3 4 5
1 2
1 3
3 4
4 5
4
2 2 2 2
3 1
3 2
3 4

Example output:
1
3
0
3

Editorial:
Let's consider what the paths passing through some vertex $$$v$$$ look like. First, root the tree arbitrarily. Let $$$u_1, u_2, \dots, u_k$$$ be the children of $$$v$$$. Then, for some color $$$x$$$, there are $$$\mathit{cnt}[u_1][x], \mathit{cnt}[u_2][x], \dots, \mathit{cnt}[u_k][x]$$$ top-level vertices in their subtrees. Top-level here means that there are no vertices of color $$$x$$$ on the path from them to $$$u_i$$$.
If the color of $$$v$$$ is not $$$x$$$, then you can combine all $$$\mathit{cnt}$$$ top-level vertices from every pair of children into paths. If the color $$$v$$$ is $$$x$$$, then all $$$\mathit{cnt}$$$ top-level vertices can only be paired with $$$v$$$. Moreover, the top-level vertices of color $$$x$$$ in subtree of $$$v$$$ is only $$$v$$$ itself now.
With these ideas, some small-to-large can be implemented. Store $$$\mathit{cnt}[v][x]$$$ for all colors $$$x$$$ such that there exist top-level vertices of this color.
In order to recalculate $$$\mathit{cnt}[v]$$$ from the values of its children, you can first calculate the sum of $$$\mathit{cnt}$$$ for each $$$x$$$, then replace $$$\mathit{cnt}[v][c_v]$$$ with $$$1$$$ (regardless of if it appeared in children or not). So that can be done by adding all values to the values of the largest child of $$$v$$$ (largest by its size of $$$\mathit{cnt}$$$, for example). During this process, you can calculate the number of paths as well.
The complexity will be $$$O(n \log^2 n)$$$ for each testcase, and that should pass freely.
There's also an idea for a faster solution. Two words: virtual trees. Basically, you can build a virtual tree of all vertices of each color. Now, there are vertices colored $$$x$$$ in it and some auxiliary vertices. The answer for that color can be calculated with some sort of dynamic programming. Similar to the first solution, for each vertex, store the number of top-level vertices of color $$$x$$$ in its subtree. All the calculations are exactly the same.
You can build all virtual trees in $$$O(n \log n)$$$ in total.