Python 高速で最小値を取れる優先度付きキュー heapq AtCoder頻出

概要

リストから最小値を高速で取り出したいときに用いる優先度付きキューがPythonにはheapqという便利なものがあるのでさくっと見直しましょう。AtCoderでも頻出内容

急いでる人向け

この記事のすべて。三行目を
q = heapq.heapify(a)
としても動かないため注意

import heapq

a = [1,2,3,4,6,7,8]
#list をheapifyでヒープ化する。
heapq.heapify(a)

#heappush で５を追加する。
heapq.heappush(a,5)

#heap内のの要素を小さい順にheappopで取り出す。
while a:
  print(heapq.heappop(a))

######実行結果######
1
2
3
4
5
6
7
8

かんたんな説明

・heap*.heapify(list)*
リストを与えるとヒープ化してくれる

・heap.heappush(heap,要素)
heapに新しい要素を追加する（listにおけるappendのheap版）

・heap.heapop(heap)
heap内の最小の値を取り出す。取り出した要素はheapから削除される。

応用

最大値を取り出す優先度付きキューを作りたい

すべての値の符号を逆転させればよい。

import heapq

a = [1,2,3,4,6,7,8]

#要素の符号を逆転
for i in range(len(a)):
  a[i] = a[i]*(-1)
heapq.heapify(a)

while a:
  #出力の符号を戻してやる
  print(heapq.heappop(a)*(-1))

######実行結果######

8
7
6
4
3
2
1

参考にしたサイト

より詳しい解説はこっち
https://docs.python.org/ja/3/library/heapq.html