New Game Plus! (Codeforces 1415E)
Tags: Greedy, constructive algorithms

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: 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: Output a single integer, the maximum score Wabbit can obtain by defeating all $$$n$$$ bosses (this value may be negative).

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$$$.

Sample Inputs:
3 0
1 1 1

5 1
-1 -2 -3 -4 5

13 2
3 1 4 1 5 -9 -2 -6 -5 -3 -5 -8 -9

Sample Outputs:
3

11

71

Solution (Python):
(n, k) = map(lambda x: int(x), input().split())
v = list(map(lambda x: int(x), input().split()))
v.sort()
k += 1
q = 0
s = 0
while len(v) > 0 and q >= 0:
    s += q
    q += v.pop()
v.append(q)
d = len(v) % k - 1
for i in range(len(v)):
    if d >= i % k:
        p = len(v) // k - i // k
    else:
        p = len(v) // k - i // k - 1
    s += v[len(v) - i - 1] * p
print(s)

Solution (Python):
(n, k) = map(lambda x: int(x), input().split())
v = list(map(lambda x: int(x), input().split()))
v.sort()
q = 0
s = 0
while len(v) > 0 and q >= 0:
    s += q
    q += v.pop()
v.append(q)
v = v[min(len(v), k + 1):]
for i in range(len(v)):
    if len(v) % (k + 1) - 1 >= i % (k + 1):
        p = len(v) // (k + 1) - i // (k + 1) + 1
    else:
        p = len(v) // (k + 1) - i // (k + 1)
    i = len(v) - i - 1
    s += v[i] * p
print(s)

Solution (Python):
(n, k) = map(lambda x: int(x), input().split())
v = list(map(lambda x: int(x), input().split()))
v.sort()
q = 0
s = 0
while len(v) > 0 and q >= 0:
    s += q
    q += v.pop()
v.append(q)
for i in range(len(v)):
    if len(v) % (k + 1) - 1 >= i % (k + 1):
        p = len(v) // (k + 1) - i // (k + 1)
    else:
        p = len(v) // (k + 1) - i // (k + 1) - 1
    i = len(v) - i - 1
    s += v[i] * p
print(s)