heapq(ヒープキュー)の使い方と、2次元リスト(リストインリスト)へのkey指定方法

はじめに

ストーン
ストーン

pythonのheapq(ヒープキュー)を使ってみようとしたところ、単純な数字だけのリストは簡単に使えるけど、自分が使いたい事例はheapqにしたいリストの中に色んなデータを持ったタプルがあるリストだったんだ(下記コード参照)。
そういったリストをヒープキューする方法があまりネット上に見つからなかったので、色んな情報をかき集めてベストではないでしょうけど、1つの方法として記録します。
私と同じような初学者ご参考になれば。

# 簡単にheapqできるリスト(いくらでもネット上で検索できる例)
list1 = [9, 3, 5, 8, 0, 1] 
heapq.heapqify(list1)
print(heapq.heapqpop(list1))
# 0
# 最小値 0が取り出せてる

# 自分がheapqで最小値を取り出したいリスト
# (x, y, flag1, flag2, distance) のタプルが入ったリスト
list2 = [(130, 90, True, False, 30), (25, 110, True, False, 99), (78, 55, False, True, 8)]
heapq.heapqify(list2)
print(heapq.heapqpop(list2)
# (25, 110, True, False, 99)
# 欲しいのはdistance(プレイヤーからの距離)の最小値8が入った(78, 55, False, True, 8)
# 先頭のxについて大小が比較され、一番小さいx = 25のデータがpopされてくる!

ヒープキューとは

ヒープ(heap) : 積み重ねたもの

キュー(queue) : 後入れ先出しというデータ管理構造

スタック(stack) : 先入れ先出しというデータ管理構造

POINT

つまり、ヒープキュー(heapq)とは、優先度付きキューと呼ばれており、爆速で(※)最小値(又は最大値)を取り出せるアルゴリズムという認識でOK!

ストーン
ストーン

私が好きなCodinGameというプログラミング学習サイトでは、ゲームAIを作製してる感覚でプログラミンを学べるのですが、ゲームのように毎ターンデータ要素入力(膨大)して、毎ターン最小値を取り出すといった用途にheapqは向いてるぞ!
A*アルゴリズム(経路探索)で用いられる各ノードのF値を計算し、F値が最小のものが最短経路となるアルゴリズムでも使えて便利!

ヒープキュー(heapq)ライブラリの使い方

import heapq ー 使用準備

何はともあれ下記の通りインポートしてください。

import heapq
# インポートしてください.
ストーン
ストーン

ちなみにheapqをテストしてみようとpythonファイル名を"heapq.py"とするとimportで自分ファイルを参照してはまるのでご注意ください。
partially initialized module 'heapq’ has no attribute 'heappush’ (most likely due to a circular import)"

heapify ー リストのヒープキュー化

次に最小値取り出しを行いたいリストをheapq化して下さい。後で出てきますがsort(), sorted()と違いリストの中身がきれいにソートされる訳ではないので(※あるルールに基づき並び替えは行われる)、heapifyしたリストの中身を見て不思議に思うかもしれませんが、ご安心ください。あくまでリストの最初の要素list[0]が最小値になっているという認識です。

# 最小値を取り出したいリストを準備
test_list = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9]

# リストをヒープキュー化 (test_list[0]を最小値に)
heapq.heapify(test_list)
print(test_list)
# [1, 1, 2, 3, 3, 8, 4, 6, 5, 5, 5, 9, 9, 7, 9]
# ※ソートではないので小さい順に並ぶ訳ではない!

heapq.heapify(x)

xはヒープ化したいリスト

heapqpop ー 最小値取り出し

heapqのメインである最小値取り出しです。heapqpopで取り出せます。1個取り出しと同時に元リストからもその1個は削除されます。

heapqpopで1個取り出した後のリストは、最小値が取り出されたことで再格付けが行われ、再びlist[0]が最小値になるように並び替えが行われています

また、heapifyでヒープ化していない場合もheapqpop出来ますが、最初の1個目のpopは最小値でも何でもないただの先頭値が出てくるだけになりますので注意。2回目以降はheapifyしていなくてもheapqpop時にヒープ化されますが。例外として空のリスト[ ] からheapqpush(後述)で挿入していけばheapifyは不要となります。

# 最小値の取り出し(test_list[0]の取り出し)(削除)
# 先に必ずheapifyしておかないと最小値出てこない.
# データ要素の削除を行いたくないときは要素0を表示する
print(test_list[0])
# 1
min_value = heapq.heappop(test_list)
print(f"min_value = {min_value}")
# min_value = 1

print(test_list)
# [1, 3, 2, 3, 5, 8, 4, 6, 5, 5, 9, 9, 9, 7]
# 1つ取り出したことでまた並び順が変わった.
# 取り出しと同時に削除も行っているので, リストの要素数も1個減っている.
# ※くどいですがtest_list[0]が最小値になってる.

# もう一回取り出し確認
min_value = heapq.heappop(test_list)
print(f"min_value = {min_value}")
# min_value = 1

print(test_list)
# [2, 3, 4, 3, 5, 8, 7, 6, 5, 5, 9, 9, 9]
# 1つ取り出したことでまたまた並び順が変わった.
# リストの要素数も1個減っている.
# ※くどいですがtest_list[0]が最小値になってる.

heapq.heapqpop(heap)

heapはヒープ化したリスト

heapqpush ー データ要素挿入

heapqpopの次はheapqpushでデータ要素の挿入です。データ挿入後のリストはヒープ化が維持されます(今挿入したデータが最小値であれば、リストの先頭に配置される)。

# heapq.heappush
# 値のpush. test_listへデータ入力と同時にリスト先頭値が最小に.
heapq.heappush(test_list, 0)
print(test_list)
# [0, 3, 2, 3, 5, 8, 4, 6, 5, 5, 9, 9, 9, 7]
# 1個増えた. 先頭が今入れた0(最小値)になった.

heapq.heappush(test_list, 7)
print(test_list)
# [0, 3, 2, 3, 5, 8, 4, 6, 5, 5, 9, 9, 9, 7, 7]
# 1個増えた.

heapq.heapqpush(heap, item)

heapはヒープ化されたリストで、値を挿入したいリスト

itemは挿入したいデータ

heapqpushpop ー データ要素挿入して取り出し

プッシュとポップを同時に行える関数もあります。同時と言いましたが厳密には関数名の通り、pushしてからのpopです。ですので今挿入したデータが最小値であればそれがそのまま取り出されてきます。

それが嫌なケースであればpopしてからpushする必要があり、その関数もあります。関数名はheapqpoppushではなく、heapqreplace 今回は割愛しました。

heapq.pushしてheapq.popしてと同じなのですがheapq.pushpopで実行した方が高速のようです(公式ドキュメントより)。

# heapq.heappushpop
# pushしてからpopする関数
min_value = heapq.heappushpop(test_list, 10)
print(f"min_value = {min_value}")
print(test_list)
# min_value = 0
# [2, 3, 4, 3, 5, 8, 7, 6, 5, 5, 9, 9, 9, 7, 10]
# プッシュ(入れて)してからポップ(出して)なので要素数変わらず.

heapq.heappushpop(heap, item)

heapはデータを挿入・取り出ししたいヒープ化されたリスト

itemはデータを挿入したい値

取り出してから挿入するheapreplace(heap, item)もあるよ

nsmallest ー n個の最小値表示

リストからn個の最小の値をリスト形式で返す関数です。heap化されてなくても大丈夫です。

n = 1の時はmin()と同じ動作をしますし、むしろmin()の方が早いそうです。また、nが小さい値を想定して作られているようでnに大きな値を指定するのは良くなさそうです(後述するnlargestも同じ)。

# heapq.nsmallest
# 最小値順にn個の値を"リスト"で返す関数
# ちなみにkey指定も出来る.
# popじゃないのでデータは削除されない
small_list = heapq.nsmallest(3, test_list)
print(small_list)
# [2, 3, 3]
print(test_list)
# [2, 3, 4, 3, 5, 8, 7, 6, 5, 5, 9, 9, 9, 7, 10]

heapq.nsmallest( n, iterablekey=None)

nは取り出したいn個の数字

iterableはリストなど

keyはリストなどの中で小さい順を判定するkeyを指定(下記実例コード参照)

nlargest ー n個の最大値表示

nsmallestのn個の最大値バージョンです。

# heapq.nlargest
# 今度はn個の大きい値を"リスト"で返す関数
# こちらもkey指定出来る.
# popじゃないのでデータ削除されないのもnsmallestと同じ
big_list = heapq.nlargest(3, test_list)
print(big_list)
# [10, 9, 9]
print(test_list)
# [2, 3, 4, 3, 5, 8, 7, 6, 5, 5, 9, 9, 9, 7, 10]

print(heapq.nlargest(3, [0, 4, 2, -1, 8, 1]))
# [8, 4, 2]
# heap化されていなくても使える

heapq.nlargest( n, iterablekey=None)

nは取り出したいn個の数字

iterableはリストなど

keyはリストなどの中で小さい順を判定するkeyを指定(下記実例コード参照)

リストインリスト(2次元リスト)でのheapqの使用例

以上でheapqの基本的な使い方は終了ですが、heap化したいリストがいつも単純な数字だけのリスト[1, 2, 3, 4,… ]とは限りません。heap化したいリストがリストインリスト(2次元リスト)の場合はどうでしょうか(タプルインリストも同じということで)。

# リストインリスト(2次元リスト)での例

# テスト1
heap_list_in_list = []
# 空のリストに最初からheappushするならheapifyは不要
heapq.heappush(heap_list_in_list, [4, "apple", 1])
heapq.heappush(heap_list_in_list, [-1, "zzzz", 999])
heapq.heappush(heap_list_in_list, [30, "banana", 100])
# heapリストの中の各要素は数字でもリスト(2次元リスト)でも良い
# 2次元リストの場合各リストの先頭で要素大小が比較される.
# key = lambda x:x[2] 等で比較要素は指定できない模様

print(heap_list_in_list)
# [[-1, 'zzzz', 999], [4, 'apple', 1], [30, 'banana', 100]]
print(heapq.heappop(heap_list_in_list))
# [-1, 'zzzz', 999]
print(heapq.heappop(heap_list_in_list))
# [4, 'apple', 1]
print(heapq.heappop(heap_list_in_list))
# [30, 'banana', 100]
# ※ 先頭要素の -1 -> 4 -> 30の順になった。
# print(heapq.heappop(heap_list_in_list))
# IndexError: index out of range
# 要素3つしかないので4回目のpopでは
# heapのリストが空で取り出せない(pop)のでエラー

結論から言うとリストインリストも使えるのですが、勝手にリストの先頭の値で大小判定されてしまいます。リストの中の3番目の要素で大小判定(並び替え)したい時ありますよね?

しかし、sorted()関数で良くあるkey指定してのソートみたいなのは出来ない様子。

ストーン
ストーン

困りましたね・・・。

なお、nsmallestやnlargestならkeyを指定出来ます。

# テスト2
heap_list_in_list = []
# 空のリストに最初からheappushするならheapifyは不要
heapq.heappush(heap_list_in_list, [4, "apple", 1])
heapq.heappush(heap_list_in_list, [-1, "zzzz", 999])
heapq.heappush(heap_list_in_list, [30, "banana", 100])

min_value = heapq.nsmallest(1, heap_list_in_list, key= lambda x:x[2])
# ラムダ関数で各要素の3番目 x[2] で比較するようkey指定
# min_value = heapq.nsmallest(1, heap_list_in_list)
# key不要(リストの先頭で比較)の場合は引数なしでOK
# ちなみにn = 1を指定した場合はmin()の方が早いらしい

print(min_value)
# [[4, 'apple', 1]]
# 要素3番目x[2]が比較されて最小値1の要素が表示された

print(heap_list_in_list)
# [[-1, 'zzzz', 999], [4, 'apple', 1], [30, 'banana', 100]]
# popではないので要素数は減っていない

max_value = heapq.nlargest(1, heap_list_in_list, key=lambda x:x[2])
print(max_value)

リストインリストでのheapq実用例

nsmallestであればkeyを指定することで思うように最小値を表示させることが出来るのは分かりました。

しかしheapqのpop, push, pushpop, replaceが使えないとheapqの売りである爆速の恩恵に預かることが出来ません。

大小比較したいkeyをheapqで指定する方法

えびちゃん

結論からいうとkeyを指定する方法はありません。
データリストの先頭要素で大小比較される仕様です。

ストーン
ストーン

では、どうするの?

先頭要素を大小比較したいデータを持ってくるようにする。

まず、考えられるのがheapq用のデータリストの仕様をheapqに合わせて、keyとなる要素を先頭に持ってくるように作り変えること。

# 元々予定していたデータ(タプル)を持ったリスト.
# (x, y, level) のタプル構成
original_list = [ (150, 70, 5),  (31, 82, 61), (36, 36, 96), (91, 48, 4)]

# heapq用にデータ(タプル)の順番を入れ替えてリストにする.
# (level, x, y)のタプル構成に作り直す(作成し直し)
heapq_list = [ (5, 150, 70), (61, 31, 82), (96, 36, 36), (4, 91, 48)]

ただし、heapq以外のプログラム部分にも修正を余儀なくされる為、一からheapqを意識して作るときはいいのでしょうが、あまり現実的ではありません。

本来のデータとkeyを持ったheapq用リストを作成(ほぼコピー)する。

現状のデータ構造(タプル)のままheapq化できるのが理想です。上記作り直しと結局似たようなことをしていますが、本来のプログラム構造を大修整するのではなく、heapq部分だけでの修正に留めることが出来ると思います。

これがベストではないと思いますが!(教えてください!)

# 元々予定していたデータ(タプル)を持ったリスト.
# (x, y, level) のタプル構成
original_list = [ (150, 70, 5),  (31, 82, 61), (36, 36, 96), (91, 48, 4)]

# heapq用にデータ(タプル)を(key, data)とした新しいリストを作る.
# (level, (x, y, level))のタプル構成を作る
heapq_list = [ (data[2], data ) for data in original_list ]
# level をkeyにするため先頭に
print(heapq_list)
# [ (5, (150, 70, 5)),  (61, (31, 82, 61)), (96, (36, 36, 96)), (4, (91, 48, 4))]
heapq.heapqify(heapq_list)
key , data = heapq.heapqpop(heapq_list)

print(data)
# (91, 48, 4) 目的のlevel最小値4が入ったタプルデータが出てくる
POINT

heapq 用にkeyを先頭に持ってきたタプル( key, original_data )のリストを用意する。

※ orignaila_data 例: ( data1, data2, key, data3 )

popもpushもheapq用リストに対して実行し、popしたデータからオリジナルに復号する。

heapq とkey指定の使用例 – ソースコード

import random
import pprint
# 2次元リストprint整形用

# (x, y)座標に出現したモンスターのレベル(Lv)
# (x, y, Lv)のタプルを管理するリストを想定
# 低いレベルのモンスターを3匹ターゲットに選びたい
# モンスターリスト
monsters = []
for i in range(10):
    x = random.randint(0, 200)
    y = random.randint(0, 100)
    level = random.randint(1, 99)
    entity = (x, y, level)
    monsters.append(entity)

print("monstersリストの表示")
pprint.pprint(monsters)
# [(55, 40, 86),
#  (164, 29, 63),
#  (196, 100, 77),
#  (31, 82, 61),
#  (121, 32, 54),
#  (91, 48, 4),
#  (36, 36, 96),
#  (126, 29, 91),
#  (26, 92, 20),
#  (93, 71, 89)]

print("")

# heapifyにkey指定が出来ないのでkey(レベル)を先頭に持ってきた専用リストを作成
# (key, (x, y, Lv))
# 例:元データ (139, 17, 42) → heap用 (42, (139, 17, 42))
monsters_for_heap = [(value[2], value) for value in monsters]
print("monsters_for_heap : heapq用にkeyを設定したリスト(heapify前)")
pprint.pprint(monsters_for_heap)
# [(86, (55, 40, 86)),
#  (63, (164, 29, 63)),
#  (77, (196, 100, 77)),
#  (61, (31, 82, 61)),
#  (54, (121, 32, 54)),
#  (4, (91, 48, 4)),
#  (96, (36, 36, 96)),
#  (91, (126, 29, 91)),
#  (20, (26, 92, 20)),
#  (89, (93, 71, 89))]
print("")

# 最大3個分
# heapq.nlagestは元々key指定できる.
big3 = heapq.nlargest(3, monsters, key= lambda x:x[2])
print("monster : level big 3")
pprint.pprint(big3, width=20)
# [(36, 36, 96),
#  (126, 29, 91),
#  (93, 71, 89)]
print("")

# 最小3個分
# heapq.nsmallestは元々key指定できる.
min3 = heapq.nsmallest(3, monsters, key= lambda x:x[2])
print("monster : level small 3")
pprint.pprint(min3, width=20)
# [(91, 48, 4),
#  (26, 92, 20),
#  (121, 32, 54)]

print("")

# heap化
heapq.heapify(monsters_for_heap)
print("monsters_for_heap =")
pprint.pprint(monsters_for_heap)
# [(4, (91, 48, 4)),
#  (20, (26, 92, 20)),
#  (77, (196, 100, 77)),
#  (61, (31, 82, 61)),
#  (54, (121, 32, 54)),
#  (86, (55, 40, 86)),
#  (96, (36, 36, 96)),
#  (91, (126, 29, 91)),
#  (63, (164, 29, 63)),
#  (89, (93, 71, 89))]
print("")

print("heappopで3回データ取り出し")
for i in range(3):
    _, monster2 = heapq.heappop(monsters_for_heap)
    # heap用に(key, 元list)のタプルを作成したので、
    # key分は捨てる(アンダーバーに代入するが使わない)
    print(monster2)
# (91, 48, 4)
# (26, 92, 20)
# (121, 32, 54)


# ちなみにsortでも出来る
print(sorted(monsters, key = lambda x: x[2]))
# [(91, 48, 4),
#  (26, 92, 20),
#  (121, 32, 54),
#  (31, 82, 61),
#  (164, 29, 63),
#  (196, 100, 77),
#  (55, 40, 86),
#  (93, 71, 89),
#  (126, 29, 91),
#  (36, 36, 96)]
ストーン
ストーン

以上です。いかがでしたでしょうか。ブログ記事内のソースコードをまとめた下記コードと参考サイトもご参照ください。

ブログ記事内のソースコードまとめ

import heapq
# ファイル名をheapq.pyとするとはまるので注意
# (ライブラリ名と同じファイル名にはしない!)
#  partially initialized module 'heapq' has no attribute 'heappush' (most likely due to a circular import)

# ヒープキューの使い方と2次元リストへの応用法(keyの指定)

# ヒープとは
# heap : 積み重ねたもの

# キューとは
# キュー(先入れ先出し)、スタック(後入れ先出し)

# ヒープ+キュー = heapqライブラリとは
# 優先度付きキュー

# heapqライブラリのオーソドックスな使い方
# 最小値を取り出したいリストを準備
test_list = [3, 1, 4, 1, 5, 9, 2, 6, 5, 3, 5, 8, 9, 7, 9]

# リストをヒープキュー化 (test_list[0]を最小値に)
heapq.heapify(test_list)
print(test_list)
# [1, 1, 2, 3, 3, 8, 4, 6, 5, 5, 5, 9, 9, 7, 9]
# ※ソートではないので小さい順に並ぶ訳ではない!

# 最小値の取り出し(test_list[0]の取り出し)(削除)
# 先に必ずheapifyしておかないと最小値出てこない.
# データ要素の削除を行いたくないときは要素0を表示する
print(test_list[0])
# 1
min_value = heapq.heappop(test_list)
print(f"min_value = {min_value}")
# min_value = 1

print(test_list)
# [1, 3, 2, 3, 5, 8, 4, 6, 5, 5, 9, 9, 9, 7]
# 1つ取り出したことでまた並び順が変わった.
# 取り出しと同時に削除も行っているので, リストの要素数も1個減っている.
# ※くどいですがtest_list[0]が最小値になってる.

# もう一回取り出し確認
min_value = heapq.heappop(test_list)
print(f"min_value = {min_value}")
# min_value = 1

print(test_list)
# [2, 3, 4, 3, 5, 8, 7, 6, 5, 5, 9, 9, 9]
# 1つ取り出したことでまたまた並び順が変わった.
# リストの要素数も1個減っている.
# ※くどいですがtest_list[0]が最小値になってる.

# heapq.heappush
# 値のpush. test_listへデータ入力と同時にリスト先頭値が最小に.
heapq.heappush(test_list, 0)
print(test_list)
# [0, 3, 2, 3, 5, 8, 4, 6, 5, 5, 9, 9, 9, 7]
# 1個増えた. 先頭が今入れた0(最小値)になった.

heapq.heappush(test_list, 7)
print(test_list)
# [0, 3, 2, 3, 5, 8, 4, 6, 5, 5, 9, 9, 9, 7, 7]
# 1個増えた.

# heapq.heappushpop
# pushしてからpopする関数
min_value = heapq.heappushpop(test_list, 10)
print(f"min_value = {min_value}")
print(test_list)
# min_value = 0
# [2, 3, 4, 3, 5, 8, 7, 6, 5, 5, 9, 9, 9, 7, 10]
# プッシュ(入れて)してからポップ(出して)なので要素数変わらず.

# heapq.nsmallest
# 最小値順にn個の値を"リスト"で返す関数
# ちなみにkey指定も出来る.
# popじゃないのでデータは削除されない
small_list = heapq.nsmallest(3, test_list)
print(small_list)
# [2, 3, 3]
print(test_list)
# [2, 3, 4, 3, 5, 8, 7, 6, 5, 5, 9, 9, 9, 7, 10]

print("")

# heapq.nlargest
# 今度はn個の大きい値を"リスト"で返す関数
# こちらもkey指定出来る.
# popじゃないのでデータ削除されないのもnsmallestと同じ
big_list = heapq.nlargest(3, test_list)
print(big_list)
# [10, 9, 9]
print(test_list)
# [2, 3, 4, 3, 5, 8, 7, 6, 5, 5, 9, 9, 9, 7, 10]
print(heapq.nlargest(3, [0, 4, 2, -1, 8, 1]))
# [8, 4, 2]
# heap化されていなくても使える
print("")

# リストインリスト(2次元リスト)での例

# テスト1
heap_list_in_list = []
# 空のリストに最初からheappushするならheapifyは不要
heapq.heappush(heap_list_in_list, [4, "apple", 1])
heapq.heappush(heap_list_in_list, [-1, "zzzz", 999])
heapq.heappush(heap_list_in_list, [30, "banana", 100])
# heapリストの中の各要素は数字でもリスト(2次元リスト)でも良い
# 2次元リストの場合各リストの先頭で要素大小が比較される.
# key = lambda x:x[2] 等で比較要素は指定できない模様

print(heap_list_in_list)
# [[-1, 'zzzz', 999], [4, 'apple', 1], [30, 'banana', 100]]
print(heapq.heappop(heap_list_in_list))
# [-1, 'zzzz', 999]
print(heapq.heappop(heap_list_in_list))
# [4, 'apple', 1]
print(heapq.heappop(heap_list_in_list))
# [30, 'banana', 100]
# ※ 先頭要素の -1 -> 4 -> 30の順になった。
# print(heapq.heappop(heap_list_in_list))
# IndexError: index out of range
# 要素3つしかないので4回目のpopでは
# heapのリストが空で取り出せない(pop)のでエラー

# テスト2
heap_list_in_list = []
# 空のリストに最初からheappushするならheapifyは不要
heapq.heappush(heap_list_in_list, [4, "apple", 1])
heapq.heappush(heap_list_in_list, [-1, "zzzz", 999])
heapq.heappush(heap_list_in_list, [30, "banana", 100])


min_value = heapq.nsmallest(1, heap_list_in_list, key= lambda x:x[2])
# ラムダ関数で各要素の3番目 x[2] で比較するようkey指定
# min_value = heapq.nsmallest(1, heap_list_in_list)
# key不要(リストの先頭で比較)の場合は引数なしでOK
# ちなみにn = 1を指定した場合はmin()の方が早いらしい

print(min_value)
# [[4, 'apple', 1]]
# 要素3番目x[2]が比較されて最小値1の要素が表示された

print(heap_list_in_list)
# [[-1, 'zzzz', 999], [4, 'apple', 1], [30, 'banana', 100]]
# popではないので要素数は減っていない

max_value = heapq.nlargest(1, heap_list_in_list, key=lambda x:x[2])
print(max_value)

参考サイト

私が書きました!
ストーン

Python 初心者
しがない会社員が日曜プログラミングを趣味として楽しんでいます。
CodinGame(https://www.codingame.com/)でゲームチックにプログラミングしながらPythonの基礎を覚えるスタイル。覚えたことをブログにしていければと。
ソースコードエディターはVisual Studio Code を愛用。
             当サイトについて