E - Amusement Park Editorial by en_translator

We may rule out the possibility of riding the attraction with \(0\) or less fun.

Therefore, the problem can be rephrased as follows.

You may choose \(K\) elements out of the sequence \(B=(1,2,\dots A_1,1,2,\dots A_2, \dots, 1,2,\dots A_N)\). What is their maximum sum?

If the length of \(B\) is \(K\), then simply its total sum is the answer.
We now assume that the length of \(B\) exceeds \(K\).

What we want to do is to sort \(B\) and choose \(K\) largest elements, but we have to devise a way.

Let us use binary search. Consider the following True/False question.

Are there \(K\) or more values in \(B\) that are at least \(m\)? \(\cdots(*)\)

\((*)\) is \(False\) if \(n=1\), and \(True\) if \(m=2 \times10^9+1\).
In addition, \((*)\) is monotonic against \(m\).
Moreover, \((*)\) can be solved in an \(O(N)\) time.

Therefore, we may do binary search against \(m\), so we can find the minimum \(m\) that satisfies \((*)\). (Let \(M\) be the minimum \(m\) we obtained.)

When choosing \(K\) elements from sequence \(B\), the following is obviously optimal:

  • choose all the elements at least \(M\)
  • choose \(M-1\) for remaining elements (until the number of chosen elements reaches up to \(M\))

Therefore, the problem has been solved.
The complexity is \(O(N \log \max(A_i))\).

last update: