Editorial for DMOPC '20 Contest 6 P4 - Land of the Carrot Trees II
Submitting an official solution before solving the problem yourself is a bannable offence.
For each query, we can root the tree and count the number of distinct elements in each subtree.
To do so quickly, we can initially store every element in a
set data structure, and do a DFS, merging the sets of all elements in the subtree.
This takes time, however, we can optimize this by merging the smaller set into the larger set.
To prove this works, note that every time an element is inserted into a new set, the new set must have at least double the size of the old set. We start with sets of size 1, so we can double this at most times. Thus the time complexity for each query takes time.
There are multiple approaches to this subtask, but we will be covering the one that extends most naturally to the final subtask.
Note that as is fixed, there are at most distinct queries, so we can compute all possible answers and answer the queries using a lookup table.
The key observation is that, when the tree is rooted at , and you want to reroot to a direct descendant , only two subgraphs are affected: the subtree rooted at before rerooting, and the subtree rooted at after rerooting. Note that these are in fact complements of each other, so for every subtree (when we are rooted at ), if we can count the number of distinct elements in its subtree and the complement of its subtree, we can compute all queries with a simple DFS.
Root the tree arbitrarily, and using the small-to-large merging trick from the first subtask, we can count the number of distinct elements in each subtree. To count the number of distinct elements in the complement, we can employ complementary counting: that is, we can find the number of elements that occur only in the subtree.
This can also be done in , by using a frequency counter and employing the small-to-large merging technique outlined in subtask 1.
For subtask 3, we process the queries offline, in DFS traversal order.
Again, reroot the tree arbitrarily, and using the observation outlined in subtask 2, compute the number of distinct elements in each subtree and its complement in .
This time, we can key a range tree by the number of subtrees that have exactly distinct elements for . Using the same observation as subtask 2, we note that we decrease the number of subtrees with elements by 1, and increase the number of subtrees with elements by 1, where is the number of distinct elements in the subtree rooted at , and is the number of distinct elements in its complement.
Finally, each query can be answered in time by querying how many subtrees have distinct elements.