Tree with Maximum Cost
Contest: Codeforces Round 527 (Div. 3)
Difficulty rating: 1900
Tags: dfs and similar, dp, trees

You are given a tree consisting exactly of $$$n$$$ vertices. Tree is a connected undirected graph with $$$n-1$$$ edges. Each vertex $$$v$$$ of this tree has a value $$$a_v$$$ assigned to it.

Let $$$dist(x, y)$$$ be the distance between the vertices $$$x$$$ and $$$y$$$. The distance between the vertices is the number of edges on the simple path between them.

Let's define the cost of the tree as the following value: firstly, let's fix some vertex of the tree. Let it be $$$v$$$. Then the cost of the tree is $$$\sum\limits_{i = 1}^{n} dist(i, v) \cdot a_i$$$.

Your task is to calculate the maximum possible cost of the tree if you can choose $$$v$$$ arbitrarily.

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

The second line of the input contains $$$n$$$ integers $$$a_1, a_2, \dots, a_n$$$ ($$$1 \le a_i \le 2 \cdot 10^5$$$), where $$$a_i$$$ is the value of the vertex $$$i$$$.

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 one integer — the maximum possible cost of the tree if you can choose any vertex as $$$v$$$.

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

Note:
Picture corresponding to the first example:

You can choose the vertex $$$3$$$ as a root, then the answer will be $$$2 \cdot 9 + 1 \cdot 4 + 0 \cdot 1 + 3 \cdot 7 + 3 \cdot 10 + 4 \cdot 1 + 4 \cdot 6 + 4 \cdot 5 = 18 + 4 + 0 + 21 + 30 + 4 + 24 + 20 = 121$$$.

In the second example tree consists only of one vertex so the answer is always $$$0$$$.

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

Example output:
121

Example input:
1
1337

Example output:
0

Editorial:
Firstly, let's calculate the answer (let it be $$$res$$$) for some fixed vertex. Let this vertex be the vertex $$$1$$$. Just run simple dfs and calculate the result using the formula from the problem statement. Also let's calculate the sum of values (let the sum in the subtree of the vertex $$$v$$$ be $$$sum_v$$$) in each subtree of the given tree if its root is the vertex $$$1$$$. It can be easily done with simple dynamic programming.
And now the magic part: let's apply the technique which is called "re-rooting" (at least we called it so). Let's maintain the correct values in subtrees at each step of our algorithm. How will values and the answer change if we will go through the edge $$$(u, v)$$$?
The following sequence of changes will change all values correctly:
1. Firstly, it can be seen that $$$res$$$ will decrease by $$$sum_v$$$ (because the distance to each vertex in this subtree will decrease by one);
2. then $$$sum_u$$$ will decrease by $$$sum_v$$$ (because we change the root of the tree) (we need this step to maintain the correct values);
3. then $$$res$$$ will increase by $$$sum_u$$$ (because the distance to each vertex in this subtree will increase by one);
4. and then $$$sum_v$$$ will increase by $$$sum_u$$$ (because we change the root of the tree) (we need this step to maintain the correct values).
So, we can recalculate all the values we need if we go through the edge. So now we can write another one dfs to try to update the answer for each vertex as a root (as the chosen vertex).