この問題を解くのにpriority queueを使う方法がある。
atcoder.jp
Pythonでpriority queueを実装するためには２つ方法があるがどちらも欠点がある。
heapqを用いた方法
- こちらを用いて実装する方が多いと思う。でもめちゃくちゃ使いづらくないですか？
- これで用意されている関数は、リストに対してin-placeで処理を施す。
- クラスが用意されていない。
from deque import PriorityQueue を用いた方法
- クラスが用意されていて1よりも扱いやすいが、2倍ぐらい遅い。
- 中身が確認できない。(中身でfor を回す等の作業ができない。)
そこで、1をベースにPriorityQueueクラスを用意した。
pushやpopをメソッドとすることで、heapqをそのまま使うよりもスッキリ見やすく実装することができる。
また、インスタンスをそのまま実行するとheapの中身が見られるようにした。
from heapq import heapify, heappop, heappush, heappushpop
class PriorityQueue:
def __init__(self, heap):
'''
heap ... list
'''
self.heap = heap
heapify(self.heap)
def push(self, item):
heappush(self.heap, item)
def pop(self):
return heappop(self.heap)
def pushpop(self, item):
return heappushpop(self.heap, item)
def __call__(self):
return self.heap
冒頭に上げた問題で使い方の具体例を示すと、こう。
import sys
read = sys.stdin.readline
def read_ints():
return list(map(int, read().split()))
X, Y, Z, K = read_ints()
A = read_ints()
B = read_ints()
C = read_ints()
A.sort(reverse=True)
B.sort(reverse=True)
C.sort(reverse=True)
from heapq import heapify, heappop, heappush, heappushpop
class PriorityQueue:
def __init__(self, heap):
'''
heap ... list
'''
self.heap = heap
heapify(self.heap)
def push(self, item):
heappush(self.heap, item)
def pop(self):
return heappop(self.heap)
def pushpop(self, item):
return heappushpop(self.heap, item)
def __call__(self):
return self.heap
heap = []
q = PriorityQueue(heap)
q.push((-(A[0] + B[0] + C[0]), 0, 0, 0))
considered = set()
ans = []
for k_th in range(1, K+1):
heap_max, i, j, k = q.pop()
ans.append(-heap_max)
for di, dj, dk in zip([1, 0, 0], [0, 1, 0], [0, 0, 1]):
i_new, j_new, k_new = i + di, j + dj, k + dk
if i_new >= X or j_new >= Y or k_new >= Z:
continue
if (i_new, j_new, k_new) in considered:
continue
considered.add((i_new, j_new, k_new))
q.push((-(A[i_new] + B[j_new] + C[k_new]), i_new, j_new, k_new))
print(*ans, sep='\n')
