Digital Wallet (Codeforces 1866D)
Tags: Greedy, dp

There are $$$N$$$ arrays, each array has $$$M$$$ positive integer elements The $$$j$$$-th element of the $$$i$$$-th array is $$$A_{i,j}$$$.Initially, Chaneka's digital wallet contains $$$0$$$ money. Given an integer $$$K$$$. Chaneka will do $$$M-K+1$$$ operations. In the $$$p$$$-th operation, Chaneka does the following procedure:   Choose any array. Let's say Chaneka chooses the $$$x$$$-th array.  Choose an index $$$y$$$ in that array such that $$$p \leq y \leq p+K-1$$$.  Add the value of $$$A_{x, y}$$$ to the total money in the wallet.  Change the value of $$$A_{x, y}$$$ into $$$0$$$. Determine the maximum total money that can be earned!

Input: The first line contains three integers $$$N$$$, $$$M$$$, and $$$K$$$ ($$$1 \leq N \leq 10$$$; $$$1 \leq M \leq 10^5$$$; $$$1 \leq K \leq \min(10, M)$$$) — the number of arrays, the size of each array, and the constant that describes the operation constraints.The $$$i$$$-th of the next $$$N$$$ lines contains $$$M$$$ integers $$$A_{i,1}, A_{i,2}, \ldots, A_{i,M}$$$ ($$$1 \leq A_{i,j} \leq 10^6$$$) — the elements of the $$$i$$$-th array.

Output: Output an integer representing the maximum total money that can be earned.

Note: In the first example, the following is a sequence of operations of one optimal strategy:  Choosing element $$$A_{1, 1}$$$ with a value of $$$10$$$.  Choosing element $$$A_{3, 2}$$$ with a value of $$$8$$$.  Choosing element $$$A_{2, 3}$$$ with a value of $$$9$$$. So the total money earned is $$$10+8+9=27$$$.In the second example, the following is a sequence of operations of one optimal strategy:  Choosing element $$$A_{3, 2}$$$ with a value of $$$8$$$.  Choosing element $$$A_{1, 2}$$$ with a value of $$$9$$$. So the total money earned is $$$8+9=17$$$.In the third example, the following is a sequence of operations of one optimal strategy:  Choosing element $$$A_{1, 3}$$$ with a value of $$$10$$$.  Choosing element $$$A_{1, 2}$$$ with a value of $$$9$$$. So the total money earned is $$$10+9=19$$$.

Sample Inputs:
3 3 1
10 4 2
8 1 9
4 8 2

3 3 2
5 9 4
1 3 1
2 8 7

3 4 3
5 9 10 1
1 3 1 5
2 5 7 2

Sample Outputs:
27

17

19

Solution (Python):
from itertools import accumulate
import sys
import io
import os
MOD = 998244353
inf = 1 << 60

def solve():
    (n, m, k) = read_int_tuple()
    A = [read_int_list() for _ in range(n)]
    (dp, np) = ([-inf] * k, [-inf] * k)
    dp[0] = 0
    for i in range(m):
        P = list(accumulate(reversed(sorted((A[j][i] for j in range(n)))), initial=0))
        for j in range(k):
            for p in range(n + 1):
                j2 = j - p + 1
                if j2 >= 0 and j2 < k:
                    np[j2] = max(np[j2], dp[j] + P[p])
        (dp, np) = (np, dp)
    print(dp[k - 1])

def main():
    T = 1
    for t in range(T):
        solve()
BUFSIZE = 8192

class FastIO(io.IOBase):
    newlines = 0

    def __init__(self, file):
        self._file = 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')

def read_int_list():
    return list(map(int, input().split()))

def read_ints_minus_one():
    return [int(x) - 1 for x in input().split()]

def read_int_tuple():
    return map(int, input().split())

def read_encode_str(d=97):
    return [ord(x) - d for x in input()]

def read_graph(n: int, m: int, d=1):
    g = [[] for _ in range(n)]
    for _ in range(m):
        (u, v) = map(int, input().split())
        g[u - d].append(v - d)
        g[v - d].append(u - d)
    return g

def read_grid(m: int):
    return [input() for _ in range(m)]

def read_int():
    return int(input())
read_str = input
if __name__ == '__main__':
    main()

Solution (Python):
from itertools import accumulate
import sys
import io
import os
MOD = 998244353
inf = 1 << 60

def solve():
    (n, m, k) = read_int_tuple()
    A = [read_int_list() for _ in range(n)]
    dp = [[-inf] * k for _ in range(m + 1)]
    dp[0][0] = 0
    for i in range(m):
        P = sorted((A[j][i] for j in range(n)), reverse=True)
        P = list(accumulate(P))
        for j in range(k):
            for p in range(n + 1):
                j2 = j - p + 1
                if j2 >= 0 and j2 < k:
                    dp[i + 1][j2] = max(dp[i + 1][j2], dp[i][j] + (P[p - 1] if p else 0))
    print(dp[m][k - 1])

def main():
    T = 1
    for t in range(T):
        solve()
BUFSIZE = 8192

class FastIO(io.IOBase):
    newlines = 0

    def __init__(self, file):
        self._file = 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')

def read_int_list():
    return list(map(int, input().split()))

def read_ints_minus_one():
    return [int(x) - 1 for x in input().split()]

def read_int_tuple():
    return map(int, input().split())

def read_encode_str(d=97):
    return [ord(x) - d for x in input()]

def read_graph(n: int, m: int, d=1):
    g = [[] for _ in range(n)]
    for _ in range(m):
        (u, v) = map(int, input().split())
        g[u - d].append(v - d)
        g[v - d].append(u - d)
    return g

def read_grid(m: int):
    return [input() for _ in range(m)]

def read_int():
    return int(input())
read_str = input
if __name__ == '__main__':
    main()

Solution (Python):
(standard_input, packages, output_together) = (1, 1, 1)
(dfs, hashing, read_from_file) = (0, 0, 0)
de = 1
if 1:
    if standard_input:
        import io, os, sys
        input = lambda : sys.stdin.readline().strip()
        import math
        inf = math.inf

        def I():
            return input()

        def II():
            return int(input())

        def MII():
            return map(int, input().split())

        def LI():
            return list(input().split())

        def LII():
            return list(map(int, input().split()))

        def LFI():
            return list(map(float, input().split()))

        def GMI():
            return map(lambda x: int(x) - 1, input().split())

        def LGMI():
            return list(map(lambda x: int(x) - 1, input().split()))
    if packages:
        from io import BytesIO, IOBase
        import random
        import os
        import bisect
        import typing
        from collections import Counter, defaultdict, deque
        from copy import deepcopy
        from functools import cmp_to_key, lru_cache, reduce
        from heapq import merge, heapify, heappop, heappush, heappushpop, nlargest, nsmallest
        from itertools import accumulate, combinations, permutations, count, product
        from operator import add, iand, ior, itemgetter, mul, xor
        from string import ascii_lowercase, ascii_uppercase, ascii_letters
        from typing import *
        BUFSIZE = 4096
    if output_together:

        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.stdout = IOWrapper(sys.stdout)
    if dfs:
        from types import GeneratorType

        def bootstrap(f, stack=[]):

            def wrappedfunc(*args, **kwargs):
                if stack:
                    return f(*args, **kwargs)
                else:
                    to = f(*args, **kwargs)
                    while True:
                        if type(to) is GeneratorType:
                            stack.append(to)
                            to = next(to)
                        else:
                            stack.pop()
                            if not stack:
                                break
                            to = stack[-1].send(to)
                    return to
            return wrappedfunc
    if hashing:
        RANDOM = random.getrandbits(20)

        class Wrapper(int):

            def __init__(self, x):
                int.__init__(x)

            def __hash__(self):
                return super(Wrapper, self).__hash__() ^ RANDOM
    if read_from_file:
        file = open('input.txt', 'r').readline().strip()[1:-1]
        fin = open(file, 'r')
        input = lambda : fin.readline().strip()
        output_file = open('output.txt', 'w')

        def fprint(*args, **kwargs):
            print(*args, **kwargs, file=output_file)
    if de:

        def debug(*args, **kwargs):
            print('\x1b[92m', end='')
            print(*args, **kwargs)
            print('\x1b[0m', end='')
(n, m, k) = MII()
grid = [LII() for _ in range(n)]
dp = [0] * (n * k)
for i in range(m - k + 1):
    cur_ma = 0
    for j in range(n):
        cur_ma = max(cur_ma, dp[j])
    for j in range(n, n * k):
        dp[j - n] = dp[j]
    for j in range(n * k):
        (x, y) = divmod(j, n)
        v = grid[y][i + x]
        (dp[j], cur_ma) = (cur_ma + v, max(cur_ma, dp[j]))
    for j in range(1, n * k):
        dp[j] = max(dp[j - 1], dp[j])
print(max(dp))