ABC 123 D - Cake 123

問題概要

D - Cake 123
X 種類、Y 種類、Z 種類の3つの種類に重みがついている。
3種類の組み合わせで重みの合計が高い順にK個出力せよ。

解法(解説AC)

X, Y, Zをソートしておき、
Priority Queueへ最も高い重みの合計を放り込んでおく。
最も高い重みのi, j, kの状態から各種類で次点に高い重みへを1つずつずらしていき、
重みの合計をPriority Queueへ追加していく。
q.push(X(i + 1) + Y + Z)
q.push(X + Y(j + 1) + Z)
q.push(X + Y + Z(k + 1))

提出コード

x, y, z, k = map(int, input().split())
a = list(map(int, input().split()))
b = list(map(int, input().split()))
c = list(map(int, input().split()))
a.sort(reverse=True)
b.sort(reverse=True)
c.sort(reverse=True)

q = []
heapq.heappush(q, (-(a[0] + b[0] + c[0]), 0, 0, 0))
kk = k
ans = []
chk = defaultdict(lambda:defaultdict(lambda:defaultdict(bool)))

while(k):
    k -= 1
    v, ai, bi, ci = heapq.heappop(q)
    ans.append(v)
    if ai + 1 < x:
        if chk[ai + 1][bi][ci] == False:
            heapq.heappush(q, (-(a[ai + 1] + b[bi] + c[ci]), ai + 1, bi, ci))
            chk[ai + 1][bi][ci] = True
    if bi + 1 < y:
        if chk[ai][bi + 1][ci] == False:
            heapq.heappush(q, (-(a[ai] + b[bi + 1] + c[ci]), ai, bi + 1, ci))
            chk[ai][bi + 1][ci] = True
    if ci + 1 < z:
        if chk[ai][bi][ci + 1] == False:
            heapq.heappush(q, (-(a[ai] + b[bi] + c[ci + 1]), ai, bi, ci + 1))
            chk[ai][bi][ci + 1] = True
for i in range(kk):
    print(-ans[i])