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):
def d_cake_123_priority_queue(X, Y, Z, K, A, B, C):
    
    import heapq
    
    a, b, c = sorted(A, reverse=True), sorted(B, reverse=True), sorted(C, reverse=True)

    def evaluation(x, y, z):
        
        return (-(a[x] + b[y] + c[z]), x, y, z)

    
    q = []
    heapq.heappush(q, evaluation(0, 0, 0))
    ans = []

    
    for _ in range(K):
        (e, x, y, z) = heapq.heappop(q)
        ans.append(-e)
        if x + 1 < X and evaluation(x + 1, y, z) not in q:
            heapq.heappush(q, evaluation(x + 1, y, z))
        if y + 1 < Y and evaluation(x, y + 1, z) not in q:
            heapq.heappush(q, evaluation(x, y + 1, z))
        if z + 1 < Z and evaluation(x, y, z + 1) not in q:
            heapq.heappush(q, evaluation(x, y, z + 1))
    return '\n'.join(map(str, ans[:K]))

X, Y, Z, K = [int(i) for i in input().split()]
A = [int(i) for i in input().split()]
B = [int(i) for i in input().split()]
C = [int(i) for i in input().split()]
print(d_cake_123_sort_modified(X, Y, Z, K, A, B, C))

def d_cake_123_binary_search(X, Y, Z, K, A, B, C):
    
    
    import bisect
    ab = sorted([e1 + e2 for e2 in B for e1 in A])  

    
    accepted, rejected = -1, 10**11
    while abs(accepted - rejected) > 1:
        mid = (accepted + rejected) // 2
        count = sum([len(ab) - bisect.bisect_left(ab, mid - e) for e in C])
        if count >= K:
            accepted = mid
        else:
            rejected = mid

    ans = []
    for e in C:
        idx = bisect.bisect_left(ab, accepted - e)
        for i in range(idx, len(ab)):
            ans.append(e + ab[i])
    ans.sort(reverse=True)
    return ' '.join(map(str, ans))

X, Y, Z, K = [int(i) for i in input().split()]
A = [int(i) for i in input().split()]
B = [int(i) for i in input().split()]
C = [int(i) for i in input().split()]
print(d_cake_123_binary_search(X, Y, Z, K, A, B, C))

def d_cake_123_binary_search():
    
    
    import sys
    import bisect
    input = sys.stdin.readline
    X, Y, Z, K = [int(i) for i in input().split()]
    A = tuple(map(int, input().split()))
    B = tuple(map(int, input().split()))
    C = tuple(map(int, input().split()))
    ab = sorted([e1 + e2 for e2 in B for e1 in A])  

    
    accepted, rejected = -1, 10**11
    while abs(accepted - rejected) > 1:
        mid = (accepted + rejected) // 2
        count = sum([len(ab) - bisect.bisect_left(ab, mid - e) for e in C])
        if count >= K:
            accepted = mid
        else:
            rejected = mid

    ans = []
    for e in C:
        idx = bisect.bisect_left(ab, accepted - e)
        ans.extend(e + _ab for _ab in ab[idx:])
    ans.sort(reverse=True)
    return '\n'.join(map(str, ans[:K]))

print(d_cake_123_binary_search())

Verdicts: Runtime Error

Wrong Answer

Accepted