Problem: The Patisserie AtCoder sells cakes with number-shaped candles. There are X, Y and Z kinds of cakes with 1-shaped, 2-shaped and 3-shaped candles, respectively. Each cake has an integer value called _deliciousness_ , as follows:    * The deliciousness of the cakes with 1-shaped candles are A_1, A_2, ..., A_X.   * The deliciousness of the cakes with 2-shaped candles are B_1, B_2, ..., B_Y.   * The deliciousness of the cakes with 3-shaped candles are C_1, C_2, ..., C_Z.  Takahashi decides to buy three cakes, one for each of the three shapes of the candles, to celebrate ABC 123.   There are X \times Y \times Z such ways to choose three cakes.   We will arrange these X \times Y \times Z ways in descending order of the sum of the deliciousness of the cakes.   Print the sums of the deliciousness of the cakes for the first, second, ..., K-th ways in this list.

Submission attempts (in order):
from collections import Counter, deque
from fractions import gcd
from functools import lru_cache
from functools import reduce
import functools
import heapq
import itertools
import math
import numpy as np
import sys

sys.setrecursionlimit(10000)
INF = float('inf')

X, Y, Z, K = list(map(int, input().split()))
A = list(sorted(map(int, input().split()), reverse=True))
B = list(sorted(map(int, input().split()), reverse=True))
C = list(sorted(map(int, input().split()), reverse=True))

heap = []
heapq.heappush(heap, (-(A[0] + B[0] + C[0]), 0, 0, 0))
history = set()
while len(history) < K:
    d, a, b, c = heapq.heappop(heap)
    if a + 1 < X:
        heapq.heappush(heap, (-(A[a + 1] + B[b] + C[c]), a + 1, b, c))
    if b + 1 < Y:
        heapq.heappush(heap, (-(A[a] + B[b + 1] + C[c]), a, b + 1, c))
    if c + 1 < Z:
        heapq.heappush(heap, (-(A[a] + B[b] + C[c + 1]), a, b, c + 1))
    if (a, b, c) not in history:
        history.add((a, b, c))
        print(-d)


import heapq
import sys

sys.setrecursionlimit(10000)
INF = float('inf')

X, Y, Z, K = list(map(int, input().split()))
A = list(sorted(map(int, input().split()), reverse=True))
B = list(sorted(map(int, input().split()), reverse=True))
C = list(sorted(map(int, input().split()), reverse=True))

heap = []
heapq.heappush(heap, (-(A[0] + B[0] + C[0]), 0, 0, 0))
history = {(0, 0, 0)}
while len(history) < K:
    d, a, b, c = heapq.heappop(heap)
    if a + 1 < X and (a + 1, b, c) not in history:
        history.add((a + 1, b, c))
        heapq.heappush(heap, (-(A[a + 1] + B[b] + C[c]), a + 1, b, c))
    if b + 1 < Y and (a, b + 1, c) not in history:
        history.add((a, b + 1, c))
        heapq.heappush(heap, (-(A[a] + B[b + 1] + C[c]), a, b + 1, c))
    if c + 1 < Z and (a, b, c + 1) not in history:
        history.add((a, b, c + 1))
        heapq.heappush(heap, (-(A[a] + B[b] + C[c + 1]), a, b, c + 1))
    print(-d)


import os
import sys

import numpy as np

if os.getenv("LOCAL"):
    sys.stdin = open("_in.txt", "r")

sys.setrecursionlimit(10 ** 9)
INF = float("inf")
IINF = 10 ** 18
MOD = 10 ** 9 + 7
# MOD = 998244353


X, Y, Z, K = list(map(int, sys.stdin.buffer.readline().split()))
A = list(map(int, sys.stdin.buffer.readline().split()))
B = list(map(int, sys.stdin.buffer.readline().split()))
C = list(map(int, sys.stdin.buffer.readline().split()))

A = np.array(A, dtype=int)
B = np.array(B, dtype=int)
C = np.array(C, dtype=int)

AB = (A[:, None] + B[None, :]).ravel()
AB.sort()
AB = AB[::-1]


l = 0
h = IINF
while abs(l - h) > 1:
    d = (l + h) // 2
    
    cnt = np.searchsorted(d - AB, C, side='right').sum()
    if cnt >= K:
        l = d
    else:
        h = d


counts = list(np.searchsorted(h - AB, C, side='right'))
ans = []
for i, cnt in enumerate(counts):
    ans += (AB + C[i])[:cnt].tolist()
ans += [l] * (K - len(ans))
ans.sort(reverse=True)
if ans:
    print(*ans, sep='\n')


Verdicts: Time Limit Exceeded

Wrong Answer

Accepted