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])