Learning to Paint (Codeforces 1942D)
Tags: Greedy, dfs and similar

Pristine Beat - Touhou⠀Elsie is learning how to paint. She has a canvas of $$$n$$$ cells numbered from $$$1$$$ to $$$n$$$ and can paint any (potentially empty) subset of cells.  Elsie has a 2D array $$$a$$$ which she will use to evaluate paintings. Let the maximal contiguous intervals of painted cells in a painting be $$$[l_1,r_1],[l_2,r_2],\ldots,[l_x,r_x]$$$. The beauty of the painting is the sum of $$$a_{l_i,r_i}$$$ over all $$$1 \le i \le x$$$. In the image above, the maximal contiguous intervals of painted cells are $$$[2,4],[6,6],[8,9]$$$ and the beauty of the painting is $$$a_{2,4}+a_{6,6}+a_{8,9}$$$.There are $$$2^n$$$ ways to paint the strip. Help Elsie find the $$$k$$$ largest possible values of the beauty of a painting she can obtain, among all these ways. Note that these $$$k$$$ values do not necessarily have to be distinct. It is guaranteed that there are at least $$$k$$$ different ways to paint the canvas. 

Input: The first line contains a single integer $$$t$$$ ($$$1 \leq t \leq 10^3$$$) — the number of test cases.The first line of each test case contains $$$2$$$ integers $$$n$$$ and $$$k$$$ ($$$1 \leq n \leq 10^3$$$, $$$1 \leq k \leq \min(2^n, 5 \cdot 10^3)$$$) — the number of cells and the number of largest values of the beauty of a painting you must find.The next $$$n$$$ lines of each test case describe $$$a$$$ where the $$$i$$$-th of which contains $$$n-i+1$$$ integers $$$a_{i,i},a_{i,i+1},\ldots,a_{i,n}$$$ ($$$-10^6 \leq a_{i,j} \leq 10^6$$$).It is guaranteed the sum of $$$n$$$ over all test cases does not exceed $$$10^3$$$ and the sum of $$$k$$$ over all test cases does not exceed $$$5 \cdot 10^3$$$.

Output: For each test case, output $$$k$$$ integers in one line: the $$$i$$$-th of them must represent the $$$i$$$-th largest value of the beauty of a painting Elsie can obtain.

Note: In the first test case, Elsie can either paint the only cell or not paint it. If she paints the only cell, the beauty of the painting is $$$-5$$$. If she chooses not to paint it, the beauty of the painting is $$$0$$$. Thus, the largest beauty she can obtain is $$$0$$$ and the second largest beauty she can obtain is $$$-5$$$. Below is an illustration of the third test case.  

Sample Inputs:
4
1 2
-5
2 4
2 -3
-1
3 8
2 4 3
1 3
5
6 20
0 -6 -3 0 -6 -2
-7 -5 -2 -3 -4
7 0 -9 -4
2 -1 1
1 -2
-6

Sample Outputs:
0 -5 
2 0 -1 -3 
7 5 4 3 3 2 1 0 
8 8 7 7 5 5 2 2 1 1 1 1 1 1 0 0 0 0 0 -1

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

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

def readInt():
    return int(input())

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

def readStr():
    return input().strip()

def solve():
    (n, k) = readInts()
    a = [[0] * i + readList() for i in range(n)]
    if n == 1:
        return sorted([0, a[0][0]], reverse=True)[:k]
    elif n == 2:
        return sorted([0, a[0][0], a[0][1], a[1][1]], reverse=True)[:k]
    else:
        dp = [[] for _ in range(n)]
        dp[0] = sorted([0, a[0][0]], reverse=True)[:k]
        dp[1] = sorted([0, a[0][0], a[0][1], a[1][1]], reverse=True)[:k]
        cnt = [0] * n
        for i in range(2, n):
            hq = []
            heapq.heappush(hq, (-a[0][i], -1))
            heapq.heappush(hq, (-a[1][i], -1))
            for j in range(2, i + 1):
                heapq.heappush(hq, (-(a[j][i] + dp[j - 2][0]), j - 2))
                cnt[j - 2] = 1
            while len(dp[i]) < k and hq:
                (val, idx) = heapq.heappop(hq)
                dp[i].append(-val)
                if idx == -1 or cnt[idx] == len(dp[idx]):
                    continue
                heapq.heappush(hq, (-(dp[idx][cnt[idx]] + a[idx + 2][i]), idx))
                cnt[idx] += 1
            dp[i].extend(dp[i - 1])
            dp[i] = sorted(dp[i], reverse=True)[:k]
        return dp[-1]
for _ in range(int(input())):
    print(*solve())

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

def main(t):
    (n, k) = map(int, input().split())
    grid = []
    for i in range(n):
        temp = list(map(int, input().split()))
        temp = [0] * i + temp
        grid.append(temp)
    nums = [[] for _ in range(n + 2)]
    accunums = [[] for _ in range(n + 2)]
    nums[0].append(0)
    nums[1].append(0)
    accunums[1].append(0)
    accunums[0].append(0)

    def judge(i, target):
        output = 0
        for j in range(i, -1, -1):
            rest = len(accunums[j]) - bisect.bisect_left(accunums[j], target - grid[j][i])
            output += rest
        return output
    for i in range(n):
        front = -10 ** 6 * (i + 1) - 10
        rear = 10 ** 6 * (i + 1) + 10
        while front < rear:
            mid = (front + rear) // 2
            if judge(i, mid) >= k:
                front = mid + 1
            else:
                rear = mid
        bound = front
        should = 0
        for j in range(i, -1, -1):
            p = len(accunums[j]) - 1
            should += p + 1
            while p >= 0 and accunums[j][p] + grid[j][i] >= bound:
                nums[i + 2].append(accunums[j][p] + grid[j][i])
                p -= 1
        should = min(should, k)
        while len(nums[i + 2]) < should:
            nums[i + 2].append(bound - 1)
        nums[i + 2].sort()
        p = len(accunums[i + 1]) - 1
        q = len(nums[i + 2]) - 1
        while (p >= 0 or q >= 0) and len(accunums[i + 2]) < k:
            if q < 0:
                accunums[i + 2].append(accunums[i + 1][p])
                p -= 1
            elif p < 0:
                accunums[i + 2].append(nums[i + 2][q])
                q -= 1
            elif accunums[i + 1][p] >= nums[i + 2][q]:
                accunums[i + 2].append(accunums[i + 1][p])
                p -= 1
            else:
                accunums[i + 2].append(nums[i + 2][q])
                q -= 1
        accunums[i + 2] = accunums[i + 2][::-1]
    print(*accunums[i + 2][::-1])
T = int(input())
t = 1
while t <= T:
    main(t)
    t += 1

Solution (Python):
import sys
from heapq import heappush, heappop
input = sys.stdin.readline
N = 1 << 31
K = 13
mask = (1 << K) - 1

def encode(a, b, c):
    v = a + N
    v <<= K
    v |= b
    v <<= K
    v |= c
    return -v

def decode(v):
    v = -v
    c = v & mask
    v >>= K
    b = v & mask
    a = v >> K
    a -= N
    return (a, b, c)

def solve():
    (n, k) = map(int, input().split())
    A = [[0] * n for _ in range(n)]
    for i in range(n):
        lst = list(map(int, input().split()))
        for (j, a) in enumerate(lst):
            A[i][i + j] = a
    dp = [[-N] * k for _ in range(n + 1)]
    dp[0][0] = 0
    for i in range(n):
        pq = []
        heappush(pq, encode(A[0][i], 0, 0))
        for j in range(1, i + 1):
            heappush(pq, encode(A[j][i] + dp[j - 1][0], j - 1, 0))
        heappush(pq, encode(dp[i][0], i, 0))
        curr = 0
        while curr < k and len(pq):
            (v, j, idx) = decode(heappop(pq))
            dp[i + 1][curr] = v
            curr += 1
            if j > 0 and idx + 1 < k and (dp[j][idx + 1] != -N):
                nv = dp[j][idx + 1]
                if j < i:
                    nv += A[j + 1][i]
                heappush(pq, encode(nv, j, idx + 1))
    return ' '.join(map(str, dp[-1]))
T = int(input())
out = [solve() for _ in range(T)]
print('\n'.join(map(str, out)))