Editorial for DMOPC '22 Contest 2 P4 - Falling Leaves
Submitting an official solution before solving the problem yourself is a bannable offence.
Let the answer for a colour be the number of sequences where is completely removed from the tree in the shortest amount of time. Consider each colour independently. We arbitrarily pick a root for the tree. This root will not be removed in the sequence of removals. Rooting the tree allows us to uniquely identify which subtrees need to be removed in an optimal removal, so we can compute the answers. Now, we want to try all possible roots which result in different optimal removals. To do this, we want to look at the subgraphs which do not contain any node of colour . If we remove all nodes with colour , this splits the graph into multiple subgraphs. For each subgraph, the answer will be the same wherever we root the tree. Therefore, we only need to try one root in each subgraph. Once we have all the answers for each possible root, we can find the answer for the whole tree.
For each colour , we try rooting the tree in every subgraph formed by removing all nodes with colour . We then compute the answer for each root in or with tree dp. Over all colours, the total amount of roots we have to try is .
Time Complexity: or
Now, say we initially rooted the tree at node and computed the answers for all colours. When we try to move the root of the tree from to a child of , notice that the only answer that changes is the answer for the colour of node , since we moved into a different subgraph. This means that at each node, we only need to update the answer for colour. Thus, we can find the answers at all roots by rerooting the tree with dfs. This involves maintaining the last root encountered for each colour and returning info up the tree to those roots.