New Game Plus! (Codeforces 1457E)
Tags: Greedy, data structures

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 os
import sys
from io import BytesIO, IOBase
BUFSIZE = 8192

class FastIO(IOBase):
    newlines = 0

    def __init__(self, file):
        self._fd = file.fileno()
        self.buffer = 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(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')
(sys.stdin, sys.stdout) = (IOWrapper(sys.stdin), IOWrapper(sys.stdout))
input = lambda : sys.stdin.readline()

def RL():
    return map(int, sys.stdin.readline().split())

def RLL():
    return list(map(int, sys.stdin.readline().split()))

def N():
    return int(input())

def print_list(l):
    print(' '.join(map(str, l)))
(n, k) = RL()
a = RLL()
a.sort()
p = list(range(n))
if k > 0:
    s = [0]
    for i in a[::-1]:
        s.append(i + s[-1])
    s = s[1:][::-1]
    delta = 0
    r = k
    now = 0
    for i in range(1, n):
        p[i] -= delta
        if p[i] == now:
            continue
        if s[i] >= 0:
            for i in range(i + 1, n):
                p[i] -= delta
            break
        delta += 1
        p[i] = now
        r -= 1
        if r == 0:
            r = k
            now += 1
res = 0
for i in range(n):
    res += p[i] * a[i]
print(res)

Solution (Python):
import sys
input = lambda : sys.stdin.readline()

def RL():
    return map(int, sys.stdin.readline().split())

def RLL():
    return list(map(int, sys.stdin.readline().split()))

def N():
    return int(input())

def print_list(l):
    print(' '.join(map(str, l)))
(n, k) = RL()
a = RLL()
a.sort()
p = list(range(n))
if k > 0:
    s = [0]
    for i in a[::-1]:
        s.append(i + s[-1])
    s = s[1:][::-1]
    delta = 0
    r = k
    now = 0
    for i in range(1, n):
        p[i] -= delta
        if p[i] == now:
            continue
        if s[i] >= 0:
            for i in range(i + 1, n):
                p[i] -= delta
            break
        delta += 1
        p[i] = now
        r -= 1
        if r == 0:
            r = k
            now += 1
res = 0
for i in range(n):
    res += p[i] * a[i]
print(res)

Solution (Python):
(n, k) = map(int, input().split())
a = sorted(list(map(int, input().split())))
p = list(range(n))
res = 0
if k > 0:
    s = [0]
    for i in a[::-1]:
        s.append(i + s[-1])
    s = s[1:][::-1]
    delta = 0
    r = k
    now = 0
    for i in range(1, n):
        p[i] -= delta
        if p[i] == now:
            continue
        if s[i] >= 0:
            for i in range(i + 1, n):
                p[i] -= delta
            break
        delta += 1
        p[i] = now
        r -= 1
        if r == 0:
            r = k
            now += 1
for i in range(n):
    res += p[i] * a[i]
print(res)