New Game Plus!
Contest: Technocup 2021 - Elimination Round 2
Difficulty rating: 2200
Tags: constructive algorithms, greedy, math

Wabbit is playing a game with $$$n$$$ bosses numbered from $$$1$$$ to $$$n$$$. The bosses can be fought in any order. Each boss needs to be defeated exactly once. There is a parameter called boss bonus which is initially $$$0$$$.

When the $$$i$$$-th boss is defeated, the current boss bonus is added to Wabbit's score, and then the value of the boss bonus increases by the point increment $$$c_i$$$. Note that $$$c_i$$$ can be negative, which means that other bosses now give fewer points.

However, Wabbit has found a glitch in the game. At any point in time, he can reset the playthrough and start a New Game Plus playthrough. This will set the current boss bonus to $$$0$$$, while all defeated bosses remain defeated. The current score is also saved and does not reset to zero after this operation. This glitch can be used at most $$$k$$$ times. He can reset after defeating any number of bosses (including before or after defeating all of them), and he also can reset the game several times in a row without defeating any boss.

Help Wabbit determine the maximum score he can obtain if he has to defeat all $$$n$$$ bosses.

Input format:
The first line of input contains two spaced integers $$$n$$$ and $$$k$$$ ($$$1 \leq n \leq 5 \cdot 10^5$$$, $$$0 \leq k \leq 5 \cdot 10^5$$$), representing the number of bosses and the number of resets allowed.

The next line of input contains $$$n$$$ spaced integers $$$c_1,c_2,\ldots,c_n$$$ ($$$-10^6 \leq c_i \leq 10^6$$$), the point increments of the $$$n$$$ bosses.

Output format:
Output a single integer, the maximum score Wabbit can obtain by defeating all $$$n$$$ bosses (this value may be negative).

Constraints:
Time limit: 2.0 s
Memory limit: 256.0 MB

Note:
In the first test case, no resets are allowed. An optimal sequence of fights would be

- Fight the first boss $$$(+0)$$$. Boss bonus becomes equal to $$$1$$$.
- Fight the second boss $$$(+1)$$$. Boss bonus becomes equal to $$$2$$$.
- Fight the third boss $$$(+2)$$$. Boss bonus becomes equal to $$$3$$$.

Thus the answer for the first test case is $$$0+1+2=3$$$.

In the second test case, it can be shown that one possible optimal sequence of fights is

- Fight the fifth boss $$$(+0)$$$. Boss bonus becomes equal to $$$5$$$.
- Fight the first boss $$$(+5)$$$. Boss bonus becomes equal to $$$4$$$.
- Fight the second boss $$$(+4)$$$. Boss bonus becomes equal to $$$2$$$.
- Fight the third boss $$$(+2)$$$. Boss bonus becomes equal to $$$-1$$$.
- Reset. Boss bonus becomes equal to $$$0$$$.
- Fight the fourth boss $$$(+0)$$$. Boss bonus becomes equal to $$$-4$$$.

Hence the answer for the second test case is $$$0+5+4+2+0=11$$$.

Example input:
3 0
1 1 1

Example output:
3

Example input:
5 1
-1 -2 -3 -4 5

Example output:
11

Editorial:
We see that each playthrough (whether it is the first one or any of the playthroughs after the reset) is completely independent of any other playthrough. Thus, we should instead think of the problem as partitioning the $$$n$$$ bosses into $$$k+1$$$ playthroughs.
Consider a playthrough with $$$x$$$ bosses which have point increments $$$a_1,a_2,\cdots,a_{x}$$$ when fought in order. Then, the number of points that we will get is $$$(x-1)a_1 + (x-2)a_2 + \ldots + (1)a_{x-1} + (0)a_{x}$$$. A simple greedy argument tell us that within a single playthrough, we should always fight the bosses in non-increasing order of point increments, i.e. $$$a_1 \geq a_2 \geq \ldots \geq a_{x}$$$.
We can visualize this with $$$k+1$$$ stacks, each representing a playthrough. Each stack contains all of the bosses that will be fought in that playthrough, and the stack is non-increasing from top to bottom, meaning we fight the bosses from top to bottom.
We now see that $$$k=0$$$ is a fairly trivial case; just fight the bosses in non-increasing order of point increments. For the rest of the tutorial, we will assume $$$k \geq 1$$$.
For simplicity, we say that a boss is in position $$$p$$$ of a playthrough if there are $$$p$$$ more bosses below this one on its playthrough stack. Notice that if we have two bosses in two different playthroughs with point increments $$$a$$$ and $$$b$$$ with positions $$$i$$$ and $$$j$$$ in their respective playthroughs, then the total points gained from the two bosses is $$$ia + jb$$$. We now see that if $$$a \geq b$$$, then $$$i \geq j$$$ and vice versa. For example, the configuration below is not optimal because swapping the $$$-5$$$ and $$$-3$$$ gives a better answer.
Therefore, all bosses in lower-numbered positions should have point increments that are less than any boss in a higher-numbered position. This means that we can place the bosses on the stacks one at a time in non-decreasing order of point increments to reach an optimal configuration.
Call a boss good if it has a non-negative point increment, and bad otherwise. Let's fix the arrangement of bad bosses and try to place the good bosses in non-decreasing order of point increments. We can see that placing a boss with point increment $$$a$$$ on a stack with height $$$h$$$ will add $$$ah$$$ to the total, so we should always pick the stack with maximum height. Thus, all of the good bosses will always end up on the same stack, and that stack will be the stack of maximum height. Call this stack the main stack; the other $$$k$$$ stacks are the side stacks.
If there exists two side stacks whose heights differ by at least $$$2$$$, then we can always move the top-most bad boss of the taller side stack to the shorter side stack and decrease the loss in points. In the example below, the $$$-10$$$ on the left stack (of height $$$5$$$) can be moved down to the top of the right stack (of height $$$3$$$).
Thus, the maximum and minimum heights of these $$$k$$$ side stacks cannot differ by more than 1.
Let the minimum height of these $$$k$$$ stacks be $$$h$$$. If we consider the bottom $$$h$$$ bosses of all $$$k+1$$$ stacks, we see that they still must have the property that all bosses in lower positions have point increments that are less than any boss in higher positions. Thus, they must consist of the $$$h(k+1)$$$ bosses with smallest point increments.
Direct Greedy
We first sort the bosses in non-decreasing order of point increments. For each possible prefix $$$P$$$ that contains only negative point increments, we take all bosses in $$$P$$$ and distribute them evenly across the $$$k+1$$$ stacks. We then take the remaining bosses and place them on the tallest of the stacks (if there are multiple stacks of the same height, the result is the same).
Performing this computation naively takes $$$O(n^2)$$$ time, but this can be sped up using precomputation of weighted prefix and suffix sums to evaluate all configurations in $$$O(n)$$$. The final time complexity is $$$O(n \log n)$$$ due to sorting.
Smarter Greedy
Let's take all $$$n$$$ point increments and place them in a "sliding" stack such that they are arranged in non-decreasing order from bottom to top. We will slide the stack across the $$$k+1$$$ playthroughs and leave the bottom point increment behind. Shown below is an example with $$$n=8$$$ and $$$k=2$$$.
Every time we slide the stack, we reduce the total by the sum of all of the point increments that we slide down by $$$1$$$ position. Thus, we should only slide it down if the sum of those is less than $$$0$$$ as that will give us a more optimal solution.
We notice that the point increments that we slide down always form a suffix of the point increments when sorted in non-decreasing order. The moment we stop sliding is when the sum of the suffix is non-negative, giving the optimal solution. Thus, we can find the optimal configuration directly by finding the right suffix and distributing the remaining negative point increments evenly.
Genius Greedy (found by K0u1e)
Sort the point increments in non-increasing order $$$a_1,a_2,\ldots,a_n$$$. Maintain a priority queue that initially contains $$$k+1$$$ zeros. We now process the point increments in non-increasing order.
To process the current point increment $$$a_i$$$, we perform the following steps in order:
• Find the largest number $$$x$$$ in the priority queue and remove it from the priority queue
• Add $$$x$$$ to our running total
• Push $$$x+a_i$$$ back into the priority queue
The answer is the final total at the end. It will be left as an exercise to the reader to figure out why this solution is indeed equivalent to the previous solution.