PROBLEM LINK:
Practice
Contest: Division 1
Contest: Division 2
Contest: Division 3
Contest: Division 4
Author: raysh07 and iceknight1093
Tester: sushil2006
Editorialist: iceknight1093
DIFFICULTY:
Medium
PREREQUISITES:
Sorting, sets
PROBLEM:
You have an array A of length N and an K coins.
Consider the following process:
- Start with \text{score} = 0 and x = 0.
- For each i from 1 to N, in order, you can:
- Pay zero coins and do nothing.
- Pay one coin (if you have it) to increase \text{score} by A_i and x by 1.
- Pay two coins (if you have them) to increase \text{score} by A_i+x and x by 1.
For each K from 1 to 2N, find the maximum possible value of \text{score}.
Here, N \leq 2\cdot 10^5.
EXPLANATION:
Let’s recall the solution to the easy version of the problem.
Define c_1 and c_2 to be the number of type 1/2 operations we perform, respectively.
If c_1 and c_2 are fixed, it’s then optimal to choose the largest c_1 + c_2 elements of A to operate on, and push the type 2 operations to the end.
This results in a score of S + c_1 + (c_1 + 1) + \ldots + (c_1 + c_2 - 1), where S is the sum of the largest c_1 + c_2 elements of A.
Using this idea, we were able to solve for a fixed value of K in linear time by fixing the value of c_1 + c_2, and noting that it’s optimal for c_2 to be as large as possible; with the largest allowed value being \min(c_1 + c_2, K - (c_1 + c_2)).
Let’s analyze the cost a bit more.
Define S_i to be the sum of the i largest elements of A.
Now, for a fixed K, if we have i = c_1 + c_2,
- The optimal value of c_2 is \min(i, K - i).
- For this value of c_2, we have c_1 = i - c_2.
In particular c_1 always equals either 2i - K or 0: specifically, we have c_1 = 0 if i \leq K-i and c_1 = 2i - K otherwise. - The score for this (K, i) combination is then S_i + (i-1) + (i-2) + \ldots + c_1.
If c_1 = 0 this cost is independent of K, while if c_1 = 2i-K it becomes
S_i + (i-1) + (i-2) + \ldots + (2i - K).
If K is fixed, our goal is to maximize this across all i \leq K.
For some small enough i (i.e. for all i satisfying 2i \leq K), the obtained score is a constant, that being S_i + (i-1) + \ldots + 1 + 0.
So, we can precompute these scores, and just take the maximum in the appropriate prefix.
This leaves only 2i \gt K to worry about.
For each of them, the best score we computed is S_i + (i-1) + (i-2) + \ldots + (2i - K).
So, if we had an array B such that B_i = S_i + (i-1) + (i-2) + \ldots + (2i-K), the answer would just be the maximum of some subarray of B.
Suppose we compute this array B for a fixed value of K.
How does it change if we had the value K+1 instead?
It’s easy to see that it doesn’t change by a lot at all: in particular, we just need to add the value (2i - K - 1) to index i.
Of course, we need to do this for all valid i (starting from the smallest i such that 2i \gt K+1).
Depending on the parity of K, this corresponds to adding either 0, 2, 4, 6, \ldots or 1, 3, 5, 7, \ldots to some subarray of B.
It might now be tempting to try and use some sort of segment tree that supports “add arithmetic progression to range”, but doing that while supporting range maximums is not so simple (related: CF 436F).
Instead, by making a couple of observations, we can solve this problem without having to use any heavy data structures at all.
First, note that the sequence we want to add to B is always non-decreasing, i.e. if i \lt j, B_i will always increase less than B_j does.
This means, if we have indices i \lt j such that B_i \leq B_j at some point of time, index i is essentially completely irrelevant: it will never be optimal in the future since B_j will always be larger.
So, suppose we store only indices that are relevant: that is, have the potential to be the maximum element either now or sometime in the future.
Let these indices be i_1, i_2, \ldots, i_m in order from left to right.
Note that in terms of current value, we’ll have B_{i_1} \gt B_{i_2} \gt \ldots \gt B_{i_m}; since if B_{i_j} \geq B_{i_{j-1}} then index i_{j-1} wouldn’t be relevant.
Thus, if we can indeed maintain only the relevant indices, the maximum value can be obtained by just looking at the value corresponding to the leftmost among them.
The question now is, how to maintain these indices?
The simplest algorithm to do this is as follows:
- Let \text{active} denote the set of currently active indices, sorted in ascending order.
- Compute the answer for the current K using the first index of the set.
- Then, when moving to K+1:
- Add index K+1 to the back of the set, if K+1 \leq N.
- Then, for each pair of adjacent indices in \text{active}, if the second index ends up with a larger value than the first when moving to K+1, delete the first index.
Repeat this process till no more deletions need to be done.
This solution, while correct, is also quite slow: for each value of K we might end up doing \mathcal{O}(N) work by iterating through all active indices, especially if the set doesn’t change much.
To optimize this solution, we can see that there’s no need to check the entire set for changes every time K increases: instead, for each adjacent pair of elements, we can do a bit of math to compute the first time at which the second will overtake the first, and only process this pair at that instant of time.
Each time an element is deleted, it will create at most one new adjacency; and the total number of deletions is \mathcal{O}(N) since each index can only be deleted once.
Each insertion/deletion takes \mathcal{O}(\log N) time if a set is used, so this is \mathcal{O}(N\log N) total to solve for all K, more than fast enough.
Let’s recap the solution.
Process K in increasing order from 1 to 2N.
For the current value of K, let B_i denote the optimal score if we perform i operations in total.
Let \text{active} be a set of indices corresponding to potential maximums of B.
Let \text{delete} be a list of vectors, where \text{delete}[x] holds a list of indices that can be deleted at time x since they’ll no longer be optimal.
Then,
- Compute the answer for the current K by using the leftmost element of \text{active}.
- Then, insert index K+1 to active, if K+1 \leq N.
This creates a new pair of adjacent elements at the end of the set of the form (y, K+1), so compute the time at which K+1 outpaces y and append y to the appropriate list. - Then, process all indices in \text{delete}[K+1]. For each of them, delete the index from \text{active} (if it still exists), and then process the new adjacency created by the deletion, once again appending to the appropriate list.
TIME COMPLEXITY:
\mathcal{O}(N\log N) per testcase.
CODE:
Editorialist's code (PyPy3)