Gotta Spend Them Tickets

Time Limit: 2000ms

Memory Limit: 256MB

K - Gotta Spend Them Tickets

Submissions Submit

Balloon world has $n$ cities with $n-1$ roads between those cities such that it's possible to move from any city to any other city, but in order to enter city $i$, you have to use a ticket of type $i$, you can't go to city $i$ if you don't have a ticket for it.
Alice has $a_i$ tickets of type $i$, she wants to go on a journey to spend all of her tickets, she wants you to tell her how many pairs of cities ($i,j$) are a good pair for her, a pair of cities ($i,j$) is good for her if the following is true, if she started at city $i$, then it's possible to go on a journey where in the end of the journey she is at city $j$ and she used all of her tickets.

Input

The first line of input contains an integer $n$($2 \le n \le 2 \cdot 10^5$), the number of cities.
The second line contains $n$ integers $a_1, a_2, ..., a_n$($ 1 \le a_i \le 10^9$), the number of tickets Alice has for city $i$.
The next $n-1$ lines each contains two integers $x_i$, $y_i$($ 1 \le x_i,y_i \le n$)meaning there's an edge between city $x_i$ and $y_i$.
It's guaranteed that the input is a tree.

Output

Print a single integer, the number of good pairs for Alice.

Notes

In the first example, the only 2 good pairs are ($2,1$) and ($2,3$). The following is an explanation for the pair ($2,1$):
Alice is at city $2$.
Alice goes to city $3$ and uses a ticket of type $3$.
Alice goes to city $2$ and uses a ticket of type $2$.
Alice goes to city $1$ and uses a ticket of type $1$.

Samples

Input	Output
3 1 1 1 1 2 2 3 Copy	2 Copy
4 6 1 2 3 1 2 1 3 1 4 Copy	10 Copy

Let's solve an easier problem, let's say that instead of needing to use a ticket to enter a city, we need to use a ticket to use a road, then a pair of cities ($i,j$) is good if for each road on the path between $i$ and $j$ Alice has an odd number of tickets for it, and for the rest of the roads she has an even number of tickets.
Returning to the original problem, let's try to find out how many ways we can distribute the tickets of each city over the roads, when we say that a city $a$ gave a ticket to a road connected to it, then we mean Alice at some point is going to use that road to travel to city $a$. We want to distribute the tickets such that there is exactly one connected path of roads where each road in it has an odd number of tickets while the rest of the roads have an even number, or such that all roads have an even number of tickets(which means Alice started at some node $i$ and finished at the same node $i$).
But also notice that if we have two cites $a$ and $b$ connected by a road and node $a$ gave $x$ tickets to that road, then node $b$ needs to give $x-1$, $x$ or $x+1$ tickets, first meaning that $a$ gave the road one extra ticket(odd number of tickets for the road), which means Alice started at the side of city $b$ and finished at the side of city $a$, the second meaning even parity, which means either we started at the side of $b$ and finished at the side of $b$, or we started and finished at the side of $a$, third is the same as the first but with opposite direction of start and finish.
To calculate the total number of such distributions we can calculate the following dp:
$dp[i][j][k][w]$.
$i$ is city $i$.
$j$ is the amount if tickets city $i$ is giving to the road above it(in the rooted tree).
$k$ is either $0$, $1$ or $2$, meaning under city $i$ there's a road going up, down or there's no road going up nor down.
$w$ is a bool meaning did we create an odd path already or not.
But since there's only 3 options for the amount of tickets we give for each city, $j$ only needs 3 values, and with that we have an $O(n)$ solution.