Labeling the Tree with Distances
Contest: Codeforces Round 856 (Div. 2)
Difficulty rating: 2400
Tags: data structures, dp, greedy, hashing, implementation, trees

You are given an unweighted tree of $$$n$$$ vertices numbered from $$$1$$$ to $$$n$$$ and a list of $$$n-1$$$ integers $$$a_1, a_2, \ldots, a_{n-1}$$$. A tree is a connected undirected graph without cycles. You will use each element of the list to label one vertex. No vertex should be labeled twice. You can label the only remaining unlabeled vertex with any integer.

A vertex $$$x$$$ is called good if it is possible to do this labeling so that for each vertex $$$i$$$, its label is the distance between $$$x$$$ and $$$i$$$. The distance between two vertices $$$s$$$ and $$$t$$$ on a tree is the minimum number of edges on a path that starts at vertex $$$s$$$ and ends at vertex $$$t$$$.

Find all good vertices.

Input format:
The first line contains one integer $$$n$$$ ($$$2\le n\le 2\cdot 10^5$$$) — the number of vertices in the tree.

The second line contains $$$n - 1$$$ integers $$$a_1,a_2,\ldots,a_{n-1}$$$ ($$$0\le a_i < n$$$) — the given list.

Then, $$$n−1$$$ lines follow. Each of them contains two integers $$$u$$$ and $$$v$$$ ($$$1\le u,v\le n$$$) denoting an edge between vertices $$$u$$$ and $$$v$$$. It is guaranteed that the edges form a tree.

Output format:
In the first line print the number of good vertices.

In the second line, print the indices of all good vertices in ascending order.

Constraints:
Time limit: 4.0 s
Memory limit: 256.0 MB

Note:
This is the tree for the first example:

And these are two possible labelings with the elements on the list so that $$$2$$$ is a good vertex (left) and $$$4$$$ is a good vertex (right).

The square below each vertex represents its label. The black squares contain the numbers which were on the given list and the only white square contains the only number which was not on the given list.

In the second example, the only good vertex is vertex $$$3$$$.

In the third example, there are no good vertices.

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

Example output:
2
2 4

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

Example output:
1
3

Editorial:
First, count the number of occurrences of each element in the list $$$a$$$. Let these numbers be $$$c_0, c_1, \ldots, c_{n-1}$$$. Then, compute the polynomial hash of the array $$$c$$$, that is $$$\displaystyle H=\sum_{i=0}^{n-1}c_i\:b^i$$$ where $$$b$$$ is the base of the hash. Because the tree is unweighted, there are only $$$n$$$ possible values to write in the unlabeled vertex (all integers between $$$0$$$ and $$$n-1$$$). Including this extra number, the hash has $$$n$$$ possibilities: $$$H+b^i$$$ for each $$$i=0,1,\ldots, n-1$$$.
Now, let's compute the same hash for each vertex of the tree. That is, for each vertex $$$x$$$ count how many vertices are at distances $$$0,1,\ldots,n-1$$$, and if these numbers are $$$d_0, d_1, \ldots, d_{n-1}$$$ compute the value $$$\displaystyle h_x=\sum_{i=0}^{n-1}d_i\:b^i$$$.
To compute these values efficiently, we will use rerooting dynamic programming: We will make two DFS. In the first one, compute the parent of each node and the hash of each node only considering the nodes in its subtree. For each vertex $$$x$$$, if we call this hash $$$\text{dp}_x$$$ then $$$\displaystyle \text{dp}_x=1+b\sum_{\text{child}\:i}\text{dp}_i$$$. In the second one (which must be performed after the first one), for each vertex $$$x$$$ compute the hash of the parent of $$$x$$$ considering the nodes which are not in the subtree of vertex $$$x$$$. If we call this hash $$$\text{dp2}_x$$$ then $$$\displaystyle \text{dp2}_x=\text{dp}_{\text{pa}_x}-b\:\text{dp}_{x}+b\:\text{dp2}_{\text{pa}_x}$$$ where $$$\text{pa}_x$$$ is the parent of $$$x$$$.
Now, observe that $$$h_x=\text{dp}_x+b\:\text{dp2}_x$$$. With these hashes, we can compute the good vertices. A vertex $$$x$$$ will be good if and only if $$$h_x=H+b^i$$$ for some $$$i=0,1,\ldots,n-1$$$. Doing a two pointers algorithm with the two sorted lists $$$h_1, h_2, \ldots, h_n$$$ and $$$H+1, H+b, H+b^2,\ldots, H+b^{n-1}$$$ is enough to find for each $$$h_x$$$ if there is a number on the other list equal to it.
To make the solution pass, it is advisable to use hashing with multiple modules or just one big modulo.
Intended complexity: $$$\mathcal{O}(n\:\text{log}\:n)$$$.