Maximum White Subtree
Contest: Codeforces Round 627 (Div. 3)
Difficulty rating: 1800
Tags: dfs and similar, dp, graphs, trees

You are given a tree consisting of $$$n$$$ vertices. A tree is a connected undirected graph with $$$n-1$$$ edges. Each vertex $$$v$$$ of this tree has a color assigned to it ($$$a_v = 1$$$ if the vertex $$$v$$$ is white and $$$0$$$ if the vertex $$$v$$$ is black).

You have to solve the following problem for each vertex $$$v$$$: what is the maximum difference between the number of white and the number of black vertices you can obtain if you choose some subtree of the given tree that contains the vertex $$$v$$$? The subtree of the tree is the connected subgraph of the given tree. More formally, if you choose the subtree that contains $$$cnt_w$$$ white vertices and $$$cnt_b$$$ black vertices, you have to maximize $$$cnt_w - cnt_b$$$.

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

The second line of the input contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$0 \le a_i \le 1$$$), where $$$a_i$$$ is the color of the $$$i$$$-th vertex.

Each of the next $$$n-1$$$ lines describes an edge of the tree. Edge $$$i$$$ is denoted by two integers $$$u_i$$$ and $$$v_i$$$, the labels of vertices it connects $$$(1 \le u_i, v_i \le n, u_i \ne v_i$$$).

It is guaranteed that the given edges form a tree.

Output format:
Print $$$n$$$ integers $$$res_1, res_2, \dots, res_n$$$, where $$$res_i$$$ is the maximum possible difference between the number of white and black vertices in some subtree that contains the vertex $$$i$$$.

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

Note:
The first example is shown below:

The black vertices have bold borders.

In the second example, the best subtree for vertices $$$2, 3$$$ and $$$4$$$ are vertices $$$2, 3$$$ and $$$4$$$ correspondingly. And the best subtree for the vertex $$$1$$$ is the subtree consisting of vertices $$$1$$$ and $$$3$$$.

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

Example output:
2 2 2 2 2 1 1 0 2

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

Example output:
0 -1 1 -1

Editorial:
This problem is about the "rerooting" technique. Firstly, let's calculate the answer for some fixed root. How can we do this? Let $$$dp_v$$$ be the maximum possible difference between the number of white and black vertices in some subtree of $$$v$$$ (yes, the subtree of the rooted tree, i.e. $$$v$$$ and all its direct and indirect children) that contains the vertex $$$v$$$. We can calculate this dynamic programming by simple dfs, for the vertex $$$v$$$ it will look like this: $$$dp_v = a_v + \sum\limits_{to \in children(v)} max(0, dp_{to})$$$.
Okay, we can store the answer for the root somewhere. What's next? Let's try to change the root from the vertex $$$v$$$ to some adjacent to it vertex $$$to$$$. Which states of dynamic programming will change? Only $$$dp_v$$$ and $$$dp_{to}$$$. Firstly, we need to "remove" the child $$$to$$$ from the subtree of the vertex $$$v$$$: $$$dp_v = dp_v - max(0, dp_{to})$$$. Then we need to "attach" the vertex $$$v$$$ and make it a child of the vertex $$$to$$$: $$$dp_{to} = dp_{to} + max(0, dp_v)$$$. Then we need to run this process recursively from $$$to$$$ (store the answer, reroot the tree and so on) and when it ends we need to "rollback" our changes. Now $$$v$$$ is the root again and we can try the next child of $$$v$$$ as the root.
Time complexity: $$$O(n)$$$.