いつ使うのか
- あとから要素が追加される条件下で都度最大/最小の要素を早く取り出す場合。
- 選択肢の中から要素を選ぶにあたって条件が2つある場合。(条件を満たすまで取得できない要素がある。その中で最大/最小の値を取得する。)
何がいいのか
push・popにO(logN)
要素数・最小値取得にO(1)
どうするのか
-
import heapqする。 - 空のリストを用意する。or既存のリストを
heapify()する。 -
heappush(heap, item)でヒープに要素追加。 -
heappop(heap)でヒープから最小値の要素取得。
その他の操作
- 2と3を同時に行う場合は
heappushpop(heap, item)でヒープにpushしたあとにpopできる（こちらの方が効率良い） - popせずに最小の要素にアクセスする場合にはインデックスを0としてブラケット演算子でアクセス。
イメージ
考えること
ヒープは最小値の取得しかできないので最大値をとりたい場合は予め全要素を負にして突っ込んでおき、取得後に絶対値をとって戻す。
テンプレート
2のケースに対して
def main():
# 取得可能になるタイミング毎に選択肢をまとめる。
# 条件を変えていく。
for n in range(N):
# 選択可能になる選択肢をエンキュー(取り出す時に最大値が欲しいのでマイナスにする。)
# ヒープが空でなければ最も良い選択肢を取得、マイナスを戻す。
return
使用例
1のケース
import heapq
def main():
A, B, C = map(int, input().split())
K = int(input())
heap = [-A, -B, -C]
heapq.heapify(heap)
for _ in range(K):
max = heapq.heappop(heap)
heapq.heappush(heap, max * 2)
sum = 0
for i in heap:
sum += i
print(-sum)
2のケース
import heapq
def main():
N = int(input())
# day日目に選択可能になる選択肢をまとめて管理しておく。
tasks = [[] for _ in range(N + 1)]
for _ in range(N):
a, b = map(int, input().split())
tasks[a].append(b)
heap = []
points = 0
for day in range(1, N + 1):
# 選択可能になる選択肢をエンキュー(取り出す時に最大値が欲しいのでマイナスにする。)
for i in tasks[day]:
heapq.heappush(heap, -i)
# ヒープが空でなければ最も良い選択肢を取得、マイナスを戻す。
if len(heap) > 0:
points += abs(heapq.heappop(heap))
print(points)
return
上記の類題
選択する条件は2つ。
- M日後までに報酬を得ること → M日目から遡っていくと選択できるタスクが増えていく。
- その中で最大の値を取得すること
※ 2条件から選択する場合、厳しい方の条件から見ていくと良い。
import heapq
def main():
N, M = map(int, input().split())
# day日目に選択可能になる選択肢をまとめて管理しておく。
jobs = [[] for _ in range(M + 1)]
for _ in range(N):
a, b = map(int, input().split())
if a > M:
continue
jobs[a].append(b)
heap = []
salary = 0
# M日目から遡っていく。
for day in range(1, M + 1):
# 選択可能になる選択肢をエンキュー(取り出す時に最大値が欲しいのでマイナスにする。)
for i in jobs[day]:
heapq.heappush(heap, -i)
# ヒープが空でなければ最も良い選択肢を取得、マイナスを戻す。
if len(heap) > 0:
salary += abs(heapq.heappop(heap))
print(salary)
return