Black, White and Grey Tree
Contest: VK Cup 2019-2020 - Final Round (Engine)
Difficulty rating: 3000

You are given a tree with each vertex coloured white, black or grey. You can remove elements from the tree by selecting a subset of vertices in a single connected component and removing them and their adjacent edges from the graph. The only restriction is that you are not allowed to select a subset containing a white and a black vertex at once.

What is the minimum number of removals necessary to remove all vertices from the tree?

Input format:
Each test contains multiple test cases. The first line contains an integer $$$t$$$ ($$$1 \le t \le 100\,000$$$), denoting the number of test cases, followed by a description of the test cases.

The first line of each test case contains an integer $$$n$$$ ($$$1 \le n \le 200\,000$$$): the number of vertices in the tree.

The second line of each test case contains $$$n$$$ integers $$$a_v$$$ ($$$0 \le a_v \le 2$$$): colours of vertices. Gray vertices have $$$a_v=0$$$, white have $$$a_v=1$$$, black have $$$a_v=2$$$.

Each of the next $$$n-1$$$ lines contains two integers $$$u, v$$$ ($$$1 \le u, v \le n$$$): tree edges.

The sum of all $$$n$$$ throughout the test is guaranteed to not exceed $$$200\,000$$$.

Output format:
For each test case, print one integer: the minimum number of operations to solve the problem.

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

Note:
In the first test case, both vertices are white, so you can remove them at the same time.

In the second test case, three operations are enough. First, we need to remove both black vertices (2 and 4), then separately remove vertices 1 and 3. We can't remove them together because they end up in different connectivity components after vertex 2 is removed.

In the third test case, we can remove vertices 1, 2, 3, 4 at the same time, because three of them are white and one is grey. After that, we can remove vertex 5.

In the fourth test case, three operations are enough. One of the ways to solve the problem is to remove all black vertices at once, then remove white vertex 7, and finally, remove connected white vertices 1 and 3.

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

Example output:
1
3
2
3

Editorial:
Let's solve the task step by step:
1. Suppose that tree is a bamboo without the grey vertices. Such a tree can be viewed as an array of colors $$$1$$$ and $$$2$$$. We can see that if there are two adjacent vertices of equal color, we can always delete them together in one operation. We can merge adjacent vertices of the same color, and get an array of colors $$$b_1, b_2, \ldots, b_k$$$, such that $$$b_i \in \{1, 2\}; b_i \ne b_{i+1}$$$. Such an array can be defined by two numbers — $$$b_1, k$$$.We can see that such an array $$$b$$$ of length $$$k$$$ can not be deleted in less than $$$\lfloor {k \over 2} \rfloor + 1$$$ removals. It can be proved by induction. Also, you can delete all elements in this number of removals by deleting opposite leaves (after the first removal opposite leaves will have the same color).
2. Let's solve the task for a general tree without grey vertices. Let's assign the edge $$$uv$$$ with weight $$$0$$$ if $$$a_u=a_v$$$, and $$$1$$$ otherwise. Let's find the longest path (diameter) in this weighted tree, and let it be the vertices $$$v_1, v_2, \ldots, v_m$$$. We can see this path as bamboo from the previous paragraph, and find the corresponding value $$$k$$$ for this path (it is equal to diameter + 1). It is obvious that we can't delete the tree in less than $$$\lfloor {k \over 2} \rfloor + 1$$$ removals (otherwise we would be able to delete the bamboo in a smaller number of removals). Turns out that we can delete all vertices in this number of removals. We can do the same algorithm — let's delete the opposite leaves of diameter, and also let's delete all leaves in the tree that have the same color (why not). After one such removal, our path will still be a diameter (if another path becomes the diameter, then one of its leaves should have the same color, and was going to be deleted).We can find the diameter in such a 0/1 tree in linear time, or we can solve the task even simpler. We can see that we alternate the removal of black and white vertices, and we delete all the leaves with the same color. So, we can choose the first operation (delete black or white), and at each iteration just delete all corresponding leaves. It works in linear time.
3. Let's solve the task without additional constraints. Now there are the grey vertices. How do they change the solution? Let's see at the last removal — suppose we deleted vertex $$$v$$$, which was not grey. Then we can imagine that we make $$$v$$$ the root of the tree, and paint all the grey vertices in the color of their parents. Then we have a tree without grey vertices, which we can solve. Obviously, the answer for such a colored tree is not less than the answer for the initial tree (because we can make the same removals as in a colored tree). But we can see that we can't get the smaller answer, as by coloring grey vertices we effectively removed them from the tree, and the value $$$\lfloor {k \over 2} \rfloor + 1$$$ (over subsequences of black and white vertices) hasn't changed.
So, overall, the solution is to choose the first removal (1 or 2), and alternate removals of black and white vertices. For removal of $$$c \in \{ 1, 2 \}$$$ we delete all the leaves with color $$$0$$$ or $$$c$$$. Also, we can note that the tree remains connected in this process.