[競プロ用]優先度付きキュー(heapq)まとめ

いつ使うのか

1. あとから要素が追加される条件下で都度最大/最小の要素を早く取り出す場合。
2. 選択肢の中から要素を選ぶにあたって条件が2つある場合。(条件を満たすまで取得できない要素がある。その中で最大/最小の値を取得する。)

何がいいのか

push・popにO(logN)
要素数・最小値取得にO(1)

どうするのか

1. import heapqする。
2. 空のリストを用意する。or既存のリストをheapify()する。
3. heappush(heap, item)でヒープに要素追加。
4. heappop(heap)でヒープから最小値の要素取得。

その他の操作

• 2と3を同時に行う場合はheappushpop(heap, item)でヒープにpushしたあとにpopできる（こちらの方が効率良い）
• popせずに最小の要素にアクセスする場合にはインデックスを0としてブラケット演算子でアクセス。

イメージ

考えること

ヒープは最小値の取得しかできないので最大値をとりたい場合は予め全要素を負にして突っ込んでおき、取得後に絶対値をとって戻す。

テンプレート

2のケースに対して

def main():
    # 取得可能になるタイミング毎に選択肢をまとめる。

    # 条件を変えていく。
    for n in range(N):
        # 選択可能になる選択肢をエンキュー(取り出す時に最大値が欲しいのでマイナスにする。)

        # ヒープが空でなければ最も良い選択肢を取得、マイナスを戻す。

    return

使用例

1のケース

ABC096 B - Maximum Sum

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のケース

第二回 アルゴリズム実技検定 F - タスクの消化

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

上記の類題

ABC137 D - Summer Vacation

選択する条件は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

参考

• https://docs.python.org/ja/3/library/heapq.html