優先度付きキューとは

優先度付きキュー（priority queue）とは、優先度にしたがって、優先度の高いものから順番に取り出すことができるものだ。

その代表例が**ヒープキュー（heapq）**である。公式ドキュメントはこちら。（https://docs.python.org/ja/3/library/heapq.html）

特徴は大きく分けて2つある。

最小値、最大値を素早く取り出すことができる
値の挿入順は気にしなくても良い

主に使うメソッドは次のとおりだ。

heapq.heapify(heap, item)：優先度付きキューの作成
heapq.heappop(heap)：最小値の取り出し
heapq.heapush(heap, item)：キューへの挿入

heapq.heappop(heap)、heapq.heapush(heap, item)の計算量はどちらも$O(logN)$なので、通常のリストでの操作よりも少ない計算量で実装できる。

# heapq
import heapq

numbers = [3, 5, 1, 6, 2, 4]

heapq.heapify(numbers) # 優先度付きキューの作成

print(numbers)

print(heapq.heappop(numbers)) # 最小値の取り出し
print(heapq.heappop(numbers))

print(numbers)

heapq.heappush(numbers, 1) # キューへの挿入
heapq.heappush(numbers, 2)

print(numbers)

print(heapq.heappop(numbers)) # 最小値の取り出し
print(heapq.heappop(numbers))

print(numbers)

>> [1, 2, 3, 6, 5, 4]
>> 1
>> 2
>> [3, 4, 5, 6]
>> [1, 3, 2, 6, 4, 5]
>> 1
>> 2
>> [3, 4, 5, 6]

例題1: ABC137 D - Summer Vacation

def main():
    import heapq

    n, m = map(int, input().split())

    jobs = [[] for _ in range(m+1)]

    for _ in range(n): # 各日のリストに各仕事の報酬を追加
        a, b = map(int, input().split())

        if a <= m:
            jobs[a].append(b)

    heap = []
    heapq.heapify(heap)
    total = 0

    for day in range(1, m+1):
        for i in jobs[day]:
            heapq.heappush(heap, -i) # 最大値を取れるように-iでpush

        if len(heap) > 0: # もし値があれば、
            total += abs(heapq.heappop(heap)) # 最大値を取得

    print(total)

if __name__ == '__main__':
    main()

例題2: ABC141 D - Powerful Discount Tickets

def main():
    import heapq

    n, m = map(int, input().split())
    alist = list(map(lambda x: int(x)*(-1), input().split()))

    heapq.heapify(alist)

    for _ in range(m):
        temp = heapq.heappop(alist)*(-1)//2
        heapq.heappush(alist, temp*(-1))

    print(sum(map(lambda x: -x, alist)))


if __name__ == '__main__':
    main()

例題3: ABC212 D - Querying Multiset

def main():
    import heapq

    q = int(input())
    sack = []
    heapq.heapify(sack)
    total = 0

    for i in range(q):
        query = list(map(int, input().split()))
        if query[0] == 1:
            heapq.heappush(sack, query[1]-total)
        elif query[0] == 2:
            total += query[1]
        else:
            ans = heapq.heappop(sack)
            print(ans+total)

if __name__ == '__main__':
    main()

例題4: 第二回アルゴリズム検定 F - タスクの消化

def main():
    import heapq

    n = int(input())

    tasks = [[] for _ in range(n+1)]

    for _ in range(n):
        a, b = map(int, input().split())

        tasks[a].append(b)

    heap = []
    heapq.heapify(heap)
    total = 0

    for day in range(1, n+1):
        for i in tasks[day]:
            heapq.heappush(heap, -i)  # 最大値の取得

        if len(heap) > 0:
            total += abs(heapq.heappop(heap))  # 最大値を加算

        print(total)


if __name__ == '__main__':
    main()

例題5: ABC173 D - Chat in a Circle

def main():
    import heapq

    n = int(input())
    alist = list(map(int, input().split()))

    alist.sort(reverse=True)

    total = 0
    heap = []

    heapq.heappush(heap, alist[0]*(-1)) # 最初に到着した人

    for i in range(1, n):
        friendly = heapq.heappop(heap)*(-1) # 輪の中で最大値を取る
        total += friendly # 心地よさに加算する
        heapq.heappush(heap, alist[i]*(-1)) # 輪に次の人を入れる
        heapq.heappush(heap, alist[i]*(-1))

    print(total)


if __name__ == '__main__':
    main()

例題6: ABC217 E - Sorting Queries

def main():
    from collections import deque
    import heapq

    q = int(input())

    que = deque()
    heap = []

    for _ in range(q):
        query = list(map(int, input().split()))
        if query[0] == 1:
            que.append(query[1]) # queに値を追加
        elif query[0] == 2:
            if len(heap) > 0: # heapに値が存在したら
                print(heapq.heappop(heap)) # heapの最小値を取り出す
            else: # heapに値がなかったら
                print(que.popleft()) # queから値を取り出す
        else:
            while que: # queの中身をすべてソートながらheapに挿入
                heapq.heappush(heap, que.pop())


if __name__ == '__main__':
    main()

例題7: ABC223 D - Restricted Permutation

heapqを用いたトポロジカルソート。

def main():
    from heapq import heappush, heappop

    n, m = map(int, input().split())
    indeg = [0] * (n+1)  # 入次数
    out = [[] for _ in range(n+1)]

    for _ in range(m):
        a, b = map(int, input().split())
        indeg[b] += 1
        out[a].append(b)

    heap = []

    # 入次数が0の数字を探す
    for i in range(1, n+1):
        # 入次数が0ならば
        if indeg[i] == 0:
            heappush(heap, i)

    ans = []

    while heap:  # すべての数字を見るまで続行
        # 辞書順で最小
        # つまり、数字が小さいものから取り出す
        now = heappop(heap)
        ans.append(now)
        # 次の数字はどこへ向かうか
        for to in out[now]:
            indeg[to] -= 1
            # 入次数が0の数字をheapに追加
            if indeg[to] == 0:
                heappush(heap, to)

    if len(ans) == n:
        print(*ans)
    else:
        print(-1)


if __name__ == '__main__':
    main()

まとめ

heapq は実装こそ難しくないので、使い慣れると最大の武器となる。
問題に行き詰まったら、heapが使えないか確認する癖をつけておこう。