New Game Plus! (Codeforces 1456C)
Tags: Greedy, binary search

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):
import heapq
import sys, os, io
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline
(n, k) = map(int, input().split())
c = list(map(int, input().split()))
c.sort(reverse=True)
s = 0
f = 0
h = []
ans = 0
for i in range(n):
    ci = c[i]
    if not f:
        ans += s
    elif n - i <= k + 1:
        break
    else:
        x = -heapq.heappop(h)
        x += ci
        ans += x
        heapq.heappush(h, -x)
        continue
    if s + ci < 0:
        f = 1
        h = [0] * (k + 1)
        h[0] = -(s + ci)
        if n - i > k + 1:
            ans += s + ci
        heapq.heapify(h)
    else:
        s += ci
print(ans)

Solution (Python):
import io
import os
input = io.BytesIO(os.read(0, os.fstat(0).st_size)).readline

def solve():
    (n, k) = map(int, input().split())
    c = sorted(map(int, input().split()))
    best = -10 ** 32
    total = sum(c)
    stepped = sum((i * x for (i, x) in enumerate(c)))
    g = k + 1
    cur = 0
    rem = total
    for (i, x) in enumerate(c):
        rem -= x
        cur += i // g * x
        best = max(best, cur + stepped + i // g * rem)
        stepped -= rem
    print(best)
t = 1
for _ in range(t):
    solve()

Solution (Python):
import sys, io, os
BUFSIZE = 8192

class FastIO(io.IOBase):
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = io.BytesIO()
        self.writable = 'x' in file.mode or 'r' not in file.mode
        self.write = self.buffer.write if self.writable else None

    def read(self):
        while True:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            if not b:
                break
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines = 0
        return self.buffer.read()

    def readline(self):
        while self.newlines == 0:
            b = os.read(self._fd, max(os.fstat(self._fd).st_size, BUFSIZE))
            self.newlines = b.count(b'\n') + (not b)
            ptr = self.buffer.tell()
            (self.buffer.seek(0, 2), self.buffer.write(b), self.buffer.seek(ptr))
        self.newlines -= 1
        return self.buffer.readline()

    def flush(self):
        if self.writable:
            os.write(self._fd, self.buffer.getvalue())
            (self.buffer.truncate(0), self.buffer.seek(0))

class IOWrapper(io.IOBase):

    def __init__(self, file):
        self.buffer = FastIO(file)
        self.flush = self.buffer.flush
        self.writable = self.buffer.writable
        self.write = lambda s: self.buffer.write(s.encode('ascii'))
        self.read = lambda : self.buffer.read().decode('ascii')
        self.readline = lambda : self.buffer.readline().decode('ascii')

def print(*args, **kwargs):
    """"""
    (sep, file) = (kwargs.pop('sep', ' '), kwargs.pop('file', sys.stdout))
    at_start = True
    for x in args:
        if not at_start:
            file.write(sep)
        file.write(str(x))
        at_start = False
    file.write(kwargs.pop('end', '\n'))
    if kwargs.pop('flush', False):
        file.flush()
(sys.stdin, sys.stdout) = (IOWrapper(sys.stdin), IOWrapper(sys.stdout))
input = lambda : sys.stdin.readline().rstrip('\r\n')
(N, K) = map(int, input().split())
k = K + 1
C = list(map(int, input().split()))
C_pos = [c for c in C if c >= 0]
C_neg = [-c for c in C if c < 0]
C_pos.sort()
C_neg.sort(reverse=True)
R = [[] for _ in range(k)]
an = 0
for (i, c_neg) in enumerate(C_neg):
    R[i % k].append(c_neg)
    an -= c_neg * (i // k)
for c_pos in C_pos:
    an += len(R[0]) * c_pos
    R[0].append(c_pos)
sum_C_pos = sum(C_pos)
idx = len(C_neg) - 1
ans = an
while idx > 0:
    m = idx % k
    if m == 0:
        sum_C_pos -= R[0][idx // k]
        idx -= 1
        continue
    a = R[m].pop()
    an -= a
    an += sum_C_pos
    sum_C_pos -= a
    if ans < an:
        ans = an
    idx -= 1
print(ans)