Editorial for COCI '23 Contest 3 #4 Restorani


Remember to use this editorial only when stuck, and not to copy-paste code from it. Please be respectful to the problem author and editorialist.
Submitting an official solution before solving the problem yourself is a bannable offence.

We can reformulate the task in the following way: given a tree of size n and two sets of nodes A and B of size m each, we need to find the length and construct the shortest path that begins and ends in node 1, alternates between nodes from set A and nodes from set B, and visits each node from both sets at least once.

Let's call a node x interesting if x \in A, x \in B or x = 1. Notice that there are \mathcal{O}(m) interesting nodes.

Let dp[i][X][Y] be minimal length of a path if we are currently in node i, we've already visited all nodes from set X \subseteq A and set Y \subseteq B. If |X| = |Y|, node i that we are currently in belongs to set Y. The transition is:

\displaystyle dp[i][X][Y] = \min_{j \in X} dp[j][X][Y \setminus \{i\}] + d[j][i],

where d[i][j] is the distance between interesting nodes i and j. Otherwise, |X| = |Y|+1 and node i belongs to set X. Transition is similar as the one mentioned. We can solve this with dynamic programming with bitmasks. The answer is then contained in dp[1][A][B], and we can construct the solution by backtracking our dp.

To calculate the minimum path length, let's consider each edge separately and determine the minimum number of times it must be traversed. This edge divides the tree into two subtrees. Let a be the number of nodes from set A in one of these subtrees, and b be the number of nodes from set B in the same subtree. If a = b = 0, this edge will never need to be traversed, and its contribution to the answer is 0. Otherwise, we must traverse this edge at least 2 \max(1, |a-b|) times. We claim that there always exists a construction that traverses each edge exactly this many times!

In the second subtask, the tree is a chain, and we can identify it with a sequence of length n. For an arbitrary sequence, let f_A(i) be the number of elements from set A in the prefix of the sequence of length i. Similarly, let f_B(i) be the number of elements from set B in the prefix of the sequence of length i. Let's call a sequence of length k positive if it satisfies f_A(k) = f_B(k) and f_A(i) \ge f_B(i) for all i \le k. Analogously, let's call a sequence of length k negative if it satisfies f_A(k) = f_B(k) and f_A(i) \le f_B(i) for all i \le k. The entire sequence can now be divided into consecutive positive and negative subsequences (you will best understand this with an example).

Observation. Any negative sequence in reverse order is positive.

We can traverse a positive sequence by pairing the i-th smallest element from set A with the i-th smallest element from set B, from the beginning of the sequence.

First, we will traverse the entire sequence from left to right (from node 1 to the last interesting node), visiting all positive subsequences. On the way back to node 1, we will traverse all negative subsequences (which, due to the opposite direction, will effectively be positive). It is easy to verify that this achieves the minimum number of passes through each edge. The complexity of this solution is \mathcal{O}(n+m).

For an arbitrary tree, similar to the second subtask, let f_A(i) be the number of elements from set A in the subtree rooted at node i (including the node itself), and let f_B(i) be the number of elements from set B in the subtree rooted at node i. Let's call the subtree rooted at node i positive if f_A(i) \ge f_B(i), and negative if f_A(i) \le f_B(i). Note that a subtree without interesting nodes is both positive and negative. We will never need to traverse such subtrees, so this case is ignored in the following text.

We solve the third subtask by traversing the tree node by node. We say we are in transition if the last interesting node we visited belonged to set A (and now we are on the way to a node from set B). Otherwise, we say we are not in transition.

Let's assume we are currently at node i. We propose the following naive algorithm:

  • If we are not in transition and i \in A, add i to the solution path, and A := A \setminus \{i\}.
  • If we are in transition and i \in B, add i to the solution path, and B := B \setminus \{i\}.
  • If we are not in transition and there is a child j of node i whose subtree is positive, move to j.
  • If we are in transition and there is a child j of node i whose subtree is negative, move to j.
  • If none of these conditions is satisfied, go to the parent of node i (or terminate the algorithm if we are at node 1).

There are many incorrect algorithms, but the key idea is to ensure that if a subtree has elements from both sets, we will pair them within that subtree (without unnecessary exiting and returning to the subtree). This ensures that we traverse each edge the minimum required number of times.

Each time we change the sets A or B, we need to recalculate f_A(i) and f_B(i) for each i. This can be done in \mathcal{O}(n), and we will change the sets \mathcal{O}(m) times. Therefore, the overall complexity of this algorithm is \mathcal{O}(nm).

In the last subtask, we cannot go node by node; instead, we need a better algorithm to determine the order of visiting interesting locations. For each node, we will maintain the optimal order of visiting interesting locations in its subtree and sets A and B of locations that need to be paired with some locations outside the subtree of that node.

Observation. We will never have |A| > 0 and |B| > 0 for any node.

Otherwise, with the same intuition as in the previous subtask, we would want to pair elements from different sets within the same subtree (to avoid the need for an additional traversal of the edge leading outside the subtree). Specifically, for a positive subtree, we will have |A| \ge 0, and for a negative one, |B| \ge 0.

First, we will calculate the required values for each child of node i, and then we will merge the traversals of the children into one traversal. Let's assume that each child has an empty traversal. We will then traverse from positive to negative subtree and vice versa, alternately visiting nodes from set A and set B (until one of these sets becomes empty). If some child j has a non-empty traversal, and for it, |A| = |B| = 0, it is sufficient to add the traversal of its subtree to the beginning or end of the overall traversal from node i.

The remaining case is when a child, for which |A| > 0 or |B| > 0 holds, has a non-empty traversal. Now we need to choose the moment when we will traverse its subtree. However, instead of choosing this when merging multiple subtrees into a larger one, during the alternating traversal of subtrees, we can "attach" the obtained path to some unvisited node. In the future, when passing through that node, we will add the attached traversal and continue the algorithm. The added traversal will leave us in the same subtree, so the overall path will not be unnecessarily extended. This way, we bypass several cases in which our algorithm repeats an edge more times than necessary.

There are multiple ways to write the solution, and to avoid a lengthy implementation, you can find a solution using a stack and a linked list in the provided code. The overall complexity of the algorithm is \mathcal{O}(n).


Comments

There are no comments at the moment.