PROBLEM LINK:
Author:Zi Song Yeoh
DIFFICULTY:
MEDIUM
PREREQUISITES:
SQRT DECOMPOSITION, LINK-CUT TREES
EXPLANATION:
For subtask $1$, we can just apply the $O(nq)$ brute force solution. For subtask $2$, since the state of VPS doesn't change for each node, we may just do a simple tree dp to compute the sum of values in each subtree to answer all queries in $O(1)$.
We'll directly solve the full problem. Let's compute the start and end time of dfs of each node. Now, we will work on the linearized tree (each node appears twice in the linearized tree). For each node, we assign a value to it which is initially $0$. Whenever we switch on the VPS of a node $u$, we increment the value of all nodes in the subtree (which is now a subsegment of the array) by $1$ and similarly if we switch it off we decrement the value of all nodes in the subtree by $1$. It is now easy to see that the nodes that will be infected if the virus starts spreading from a node $u$ is precisely the nodes in subtree of $u$ with value $0$. Thus, we only have to maintain this info in the array efficiently.
One way to do this is by sqrt decomposition. Divide the array into $\sqrt{n}$ blocks. For each block, store $add[i]$ and $cnt[i][j]$, which we will describe later. $cnt[i][j]$ will store the number of animals in the block which have the value of node equal to $j + add[i]$. With this information, each update can be done easily. The only thing that matters is adding (or subtracting) an entire block by $1$. This is where the add array comes handy. We can just increment the add array by $1$ or $-1$ instead of iterating through the whole block again. Answering each query is also the same. To find the answer for a whole block, just refer to the cnt array. Thus, this solution will work in $O(n\sqrt{n})$ time.
AUTHOR'S AND TESTER'S SOLUTIONS:
Author's solution can be found here.
Tester's solution can be found here.