Tree Operations
Contest: Codeforces Global Round 27
Difficulty rating: 2500
Tags: binary search, brute force, dfs and similar, dp, trees

One day, a turtle gives you a tree with $$$n$$$ nodes rooted at node $$$x$$$. Each node has an initial nonnegative value; the $$$i$$$-th node has starting value $$$a_i$$$.

You want to make the values of all nodes equal to $$$0$$$. To do so, you will perform a series of operations on the tree, where each operation will be performed on a certain node. Define an operation on node $$$u$$$ as choosing a single node in $$$u$$$'s subtree$$$^{\text{∗}}$$$ and incrementing or decrementing its value by $$$1$$$. The order in which operations are performed on nodes is as follows:

- For $$$1 \le i \le n$$$, the $$$i$$$-th operation will be performed on node $$$i$$$.
- For $$$i > n$$$, the $$$i$$$-th operation will be performed on the same node as operation $$$i - n$$$.

More formally, the $$$i$$$-th operation will be performed on the $$$(((i - 1) \bmod n) + 1)$$$-th node.$$$^{\text{†}}$$$

Note that you cannot skip over operations; that is, you cannot perform the $$$i$$$-th operation without first performing operations $$$1, 2, \ldots, i - 1$$$.

Find the minimum number of operations you must perform before you can make the values of all nodes equal to $$$0$$$, assuming you pick operations optimally. If it's impossible to make the values of all nodes equal to $$$0$$$ after finite operations, output $$$-1$$$.

Input format:
The first line contains a single integer $$$t$$$ ($$$1\le t\le 100$$$) — the number of test cases.

The first line of each test case contains two integers $$$n$$$ and $$$x$$$ ($$$1 \le n \le 2000$$$, $$$1 \le x \le n$$$) — the number of nodes and the root of the tree.

The second line of each test case contains $$$n$$$ integers $$$a_1, a_2, \ldots, a_n$$$ ($$$0 \le a_i \le 10^9$$$) — the starting value of each node.

Each of the next $$$n - 1$$$ lines of each test case contains two integers $$$u$$$ and $$$v$$$ ($$$1 \le u, v \le n$$$, $$$u \neq v$$$) representing an undirected edge from $$$u$$$ to $$$v$$$. It is guaranteed that the given edges form a tree.

It is guaranteed that the sum of $$$n$$$ over all test cases does not exceed $$$2000$$$.

Output format:
For each test case, output a single integer denoting the minimum amount of operations needed to make all nodes $$$0$$$. If it's impossible to make all nodes $$$0$$$, output $$$-1$$$.

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

Note:
In the first test case, you can make the following valid sequence of operations:

- For operation $$$1$$$, decrease the value of node $$$1$$$. This is valid because $$$(((1 - 1) \bmod n) + 1) = 1$$$, and node $$$1$$$ is in the subtree of node $$$1$$$.
- For operation $$$2$$$, decrease the value of node $$$2$$$. This is valid because $$$(((2 - 1) \bmod n) + 1) = 2$$$, and node $$$2$$$ is in the subtree of node $$$2$$$.
- For operation $$$3$$$, decrease the value of node $$$2$$$. This is valid because $$$(((3 - 1) \bmod n) + 1) = 1$$$, and node $$$2$$$ is in the subtree of node $$$1$$$.

Example input:
5
2 1
1 2
1 2
3 2
2 1 3
2 1
3 2
4 1
1 1 0 1
1 2
2 3
1 4
12 6
14 4 5 6 12 9 5 11 6 2 1 12
3 9
10 6
6 12
4 3
3 1
5 11
9 7
5 6
1 8
2 8
5 1
1 1
0

Example output:
3
6
5
145
0

Editorial:
The key observation to make is that if we can make all nodes equal to $$$0$$$ using $$$t$$$ operations, we can make all nodes equal to $$$0$$$ in $$$t + 2n$$$ operations by increasing all nodes by $$$1$$$ before decreasing all nodes by $$$1$$$.
As such, this motivates a binary search solution. For each $$$i$$$ from $$$0$$$ to $$$2n - 1$$$, binary search on the set of integers {$$$x \in \mathbb{Z} \mid x \bmod 2n \equiv i$$$} for the first time we can make all nodes $$$0$$$. Our answer is then the minimum value over all these binary searches. For a constant factor speedup, note that we only need to check the values of $$$i$$$ such that the parity of $$$i$$$ is equal to the parity of $$$\sum{a_i}$$$.
All that remains is to find a suitably fast check function for determining whether or not we can solve the tree using a fixed number of operations.
This can be done with dynamic programming in $$$O(n)$$$ time. Suppose we fix the number of operations as $$$x$$$. We can then calculate for each node in $$$O(1)$$$ time how many operations on its respective subtree it can perform. For a node $$$i$$$, this is $$$\lfloor \frac{x}{n} \rfloor + [((x \bmod n) + 1) < i]$$$. Define this value as $$$\text{has_i}$$$. Define $$$\text{dp}[i]$$$ as the number of operations needed to make the entire subtree of $$$i$$$ equal to $$$0$$$. $$$\text{dp}[i]$$$ is then $$$(\sum_{j\in\text{child($$$i$$$)}}\text{dp}[j]) + a_i - \text{has_i}$$$. If $$$\text{dp}[i] < 0$$$ at the end, $$$\text{dp}[i] = -\text{dp}[i] \bmod 2$$$, since we may need to waste an operation. Finally, the tree is solvable for a fixed $$$x$$$ if $$$\text{dp}[\text{root}] = 0$$$.
This solution can actually be improved to $$$O(n\log(a_i) + n^2\log(n))$$$. We will do so by improving our bounds for each binary search we perform. This improved solution also serves as a proof for why it is always possible to make all nodes equal to $$$0$$$ after finite operations.
Consider solving the easier subtask where the goal is to make all values of the tree $$$\le 0$$$. The minimum number of operations can be binary searched for directly, and the only change to the check function is changing the value of $$$\text{dp}[i]$$$ for when $$$\text{dp}[i]$$$ $$$\le 0$$$ to simply be $$$0$$$ instead. As such, this subtask can be solved in $$$O(n\log(a_i))$$$, and is clearly a lower bound on the answer. Let's denote this number of operations as $$$t$$$.
Claim: The tree is solvable within $$$t + n^2 + 2n + 1$$$ operations.
Proof: After using $$$t$$$ operations, we can reduce all the values in the tree into either $$$1$$$ or $$$0$$$ based on the parity of extra operations used. We want to turn all the nodes into the same number. Consider running through $$$n$$$ operations with each node applying the operation to itself. Note that the parity of all nodes flip. We will fix the parity of $$$1$$$ node per $$$n$$$ operations (that is, while all other nodes change parity, one node's parity will stay the same). We can accomplish this by having the root node change the parity of the node who's parity we want to maintain, while all other nodes apply the operation to itself.
Now, after $$$n^2$$$ operations, all nodes (except for possibly the root) will have the same parity. We can fix the parity of the root in at worst $$$n + 1$$$ extra operations. Have all nodes toggle their own parity, but have the root toggle the parity of the node that comes directly after it. Now, that node can just fix itself, and all nodes have the same parity. If all the nodes are equal to $$$1$$$, we simply use $$$n$$$ operations to toggle the parity and make all the nodes equal to $$$0$$$.
With this bound of $$$t + n^2 + 2n + 1$$$, we can simply adjust the upper and lower bounds of our binary search for the $$$O(n^2\log(a_i))$$$ solution, resulting in a $$$O(n\log(a_i) + n^2\log(n))$$$ solution.