P16371 [IATI 2026] Infection

Athena is a grown-up now and is already one of the VIP (Very Important Party-goer)-s in the country of Slopia. As Slopia has had very efficient designers, the country has $N$ cities connected with $N-1$ bidirectional roads $(U_j, V_j)$ between them, such that it's possible to travel between any $2$ cities using only the roads. However, last month the government in Slopia found out that there is a highly contagious infection going around and is now concerned with the well-being of the VIP-s. To that end, the government tracked down the precise plans of all $M$ party-goers: the day (since the beginning of the universe) they will start partying $D_i$, the city $S_i$ they will start in, and the city $T_i$ they will end in. Since the trendy way to party is to party in a different city every night, each party-goer $i$ will arrive in Slopia on the evening of day $D_i$ in city $S_i$ and party there that night, then they will continue by going from city to city, travelling one road per day, following the shortest path to $T_i$ and attending one party every night in the city they are currently at. After partying in $T_i$, they will take a flight out of Slopia on the next morning and this will end their party-tour. Additionally the Slopia government has acquired information about whether each party-goer had the infection at the beginning of their party-tour, denoted by $I_i$. Since parties are a very social endeavour, if someone has the infection at a party in some city, everyone else who is partying in the same city on the same night gets the infection. Since the infection has no cure yet, each infected person stays infected throughout the remainder of their party-tour. The Slopia government wants to calculate how many nights each party-goer will spend infected during their party-tour. It is up to you to help them by implementing an efficient solution to this problem. ### Implementation details You have to implement the function $\verb|solve|$: ```cpp std::vector solve( std::vector R, std::vector D, std::vector S, std::vector T, std::vector I ); ``` - $R$: vector of $N-1$ pairs of integers -- the $2$ cities $U_j$ and $V_j$ connected by road $j$; - $D$: vector of $M$ integers -- the start day of party-goer $i$; - $S$: vector of $M$ integers -- the start city of party-goer $i$; - $T$: vector of $M$ integers -- the final city of party-goer $i$; - $I$: vector of $M$ booleans -- the initial infection status of party-goer $i$, denoted as a boolean with $1$ meaning infected and $0$ meaning non-infected. This function will be called once per test and must return a vector of $M$ integers -- the number of nights party-goer $i$ will spend infected.

Input format: - $N \ M$ - $U_{0} \ V_{0}$ - $U_{1} \ V_{1}$ - $\cdots$ - $U_{N-2} \ V_{N-2}$ - $D_0 \ S_0 \ T_0 \ I_0$ - $\cdots$ - $D_{M-1} \ S_{M-1} \ T_{M-1} \ I_{M-1}$

Output format: - $Ans_0$ - $Ans_1$ - $\cdots$ - $Ans_{M-1}$

### Slopia Infrastructure :::align{center}  ::: The party-routes of all VIP-s are shown in the following table, where the days on which each person had the infection are in $\textbf{bold}$. | -- | **0** | **1** | **2** | **3** | **4** | | :---: | :---: | :---: | :---: | :---: | :---: | | 0 | -- | -- | -- | -- | -- | | 1 | **0** | -- | -- | -- | -- | | 2 | **1** | -- | -- | -- | -- | | 3 | **2** | 5 | -- | -- | -- | | 4 | **3** | 4 | 5 | -- | -- | | 5 | **4** | 3 | **4** | -- | -- | | 6 | **5** | 6 | **3** | -- | -- | | 7 | -- | 7 | **6** | -- | -- | | 8 | -- | -- | **7** | **7** | -- | | 9 | -- | -- | -- | **6** | -- | | 10 | -- | -- | -- | **8** | 0 | | 11 | -- | -- | -- | -- | 1 | ### Constraints - $1 \leq N, M \leq 150000$ - $0 \leq D_i \leq 10^{12}$ - $0 \leq U_j, V_j, S_i, T_i < N$ - $I_i \in \{0, 1\}$ ### Subtasks | **Subtask** | **Points** | **Required subtasks** | **$N, M$** | **Other constraints** | | :---: | :---: | :---: | :---: | :---: | | $0$ | $0$ | $--$ | $--$ | Sample test. | | $1$ | $5$ | $--$ | $\leq 100$ | $D_i = 0$ | | $2$ | $6$ | $1$ | $\leq 5000$ | $D_i = 0$ | | $3$ | $7$ | $0 - 2$ | $\leq 5000$ | $--$ | | $4$ | $5$ | $--$ | $\leq 150000$ | $D_i = 10^6 \times i$ | | $5$ | $20$ | $--$ | $\leq 150000$ | $(U_j, V_j) = (j, j + 1)$
$D_i = 0$ | | $6$ | $19$ | $5$ | $\leq 150000$ | $(U_j, V_j) = (j, j + 1)$ | | $7$ | $18$ | $0 - 3$ | $\leq 100000$ | $--$ | | $8$ | $20$ | $0 - 7$ | $\leq 150000$ | $--$ | The points for a given subtask are obtained only if all the tests for it and the required subtasks are successfully passed.