Two Colors
Contest: Kotlin Heroes: Episode 9 (Unrated, T-Shirts + Prizes!)
Difficulty rating: 2900
Tags: *special

You are given a tree consisting of $$$n$$$ vertices. Some vertices of the tree are red, all other vertices are blue.

Each edge of the tree has a positive weight. Let's define $$$d(x, y)$$$ as the distance between the vertices $$$x$$$ and $$$y$$$, i. e. the sum of weights of edges belonging to the simple path between $$$x$$$ and $$$y$$$.

For each vertex, you have to choose an integer $$$v_i$$$. These integers should meet the following constraint: for every blue vertex $$$b$$$ and every red vertex $$$r$$$, $$$d(b, r) \ge v_b + v_r$$$.

You have to maximize the value of $$$\sum \limits_{i=1}^{n} v_i$$$.

Note that the values of $$$v_i$$$ are not necessarily positive.

Input format:
The first line contains one integer $$$n$$$ ($$$2 \le n \le 3 \cdot 10^5$$$).

The second line contains $$$n$$$ integers $$$c_1, c_2, \dots, c_n$$$ ($$$0 \le c_i \le 1$$$). If the $$$i$$$-th vertex is red, $$$c_i = 1$$$, otherwise $$$c_i = 0$$$.

Then $$$n-1$$$ lines follow. Each line contains three integers $$$x_i$$$, $$$y_i$$$ and $$$w_i$$$ ($$$1 \le x_i, y_i \le n$$$; $$$1 \le w_i \le 10^6$$$; $$$x_i \ne y_i$$$) denoting an edge between the vertices $$$x_i$$$ and $$$y_i$$$ which has weight $$$w_i$$$. These edges form a tree.

Output format:
If the value of $$$\sum \limits_{i=1}^{n} v_i$$$ can be as big as possible, print Infinity. Otherwise, print one integer — the maximum possible value of $$$\sum \limits_{i=1}^{n} v_i$$$ you can get.

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

Note:
In the first example, you can assign $$$v_1 = 120, v_2 = 20, v_3 = 80, v_4 = 130$$$.

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

Example output:
350

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

Example output:
203

Editorial:
It's quite easy to see that if the number of red vertices is not equal to the number of blue vertices, the answer can be as large as you want. To prove it rigorously, we can notice that if we obtained some answer, by subtracting $$$1$$$ from vertices of one color and adding $$$1$$$ to vertices of the other color, we can increase it.
Okay, now let's deal with the case when the number of red vertices is equal to the number of blue vertices. There is a way to handle it with some mathematical transformations related to linear programming duality, but the model approach is different. Basically, the fact that the number of red vertices is the same as the number of blue vertices is a hint that this problem is related to matchings. You can try recalling solutions to different matching-based problems and noticing anything similar in the statement.
The constraints given in the statement are exactly the constraints for the potentials in the Hungarian algorithm for the weighted matching problem (also known as the assignment problem). So, we have to solve the assignment problem on the matrix where the integers in the matrix are equal to distances in the tree between the corresponding vertices. The Hungarian algorithm states that the sum of potentials is equal to the minimum weight of the perfect matching; so we need to somehow solve the following problem:
You are given a weighted tree, where each vertex has one of two colors, red or blue. The number of red vertices is equal to the number of blue vertices. Match the vertices of different colors so that the sum of distances between vertices in each pair is the minimum possible.
We have to solve it in something like $$$O(n)$$$ or $$$O(n \log n)$$$. Fortunately, it is not that difficult.
Let's root the tree at an arbitrary vertex, and for each vertex, calculate the number of red and blue vertices in its subtree. The difference between these two numbers (its absolute value) corresponds to the minimum number of times the edge leading into this subtree should be counted in the answer (i. e. if there are $$$3$$$ red and $$$5$$$ blue vertices in the subtree, at least $$$2$$$ pairs of vertices consist of a vertex in the subtree and an outside vertex, so the edge leading into the subtree is counted at least twice in the answer).
Okay, we got a lower bound for the answer. Let's show that it's reachable. I think there are ways to do this without using flows, but our proof is based on them.
If we model our problem using minimum cost flows, we have to build the following flow network: each edge of the tree transforms into an undirected edge with infinite capacity and cost equal to the weight of the edge. Then, a source is added, and for each red vertex in the tree, we send a directed edge from the source to it with capacity equal to $$$1$$$. Then, we add a sink and from each blue vertex, add a directed edge from it to the sink with capacity equal to $$$1$$$.
Let's transform our lower bound for the answer into a valid maximum flow for this network. For each vertex, if the number of red vertices in its subtree is $$$x$$$ and the number of blue vertices in its subtree is $$$y$$$, let the edge from it to its parent have flow equal to $$$x - y$$$ (if $$$y > x$$$, then $$$y - x$$$ units of flow are transferred from the parent into the subtree). It's quite easy to see that if we also add $$$1$$$ unit of flow on all edges that are incident to the source or the sink, we will get a valid flow.
So, our lower bound for the answer is reachable, and we can calculate it in $$$O(n)$$$.