AF - 4.03. heapq・bisect 解説 /

実行時間制限: 0 msec / メモリ制限: 0 KiB

前のページ | 次のページ

キーポイント

  • heapq は優先度付きキューを扱うモジュールで、要素追加と最小値や最大値の取得が可能
  • bisect モジュールの関数を利用することでソート済みのリストに対し、一定値以下の要素の個数や値の検索を行うことができる

heapq

heapq は優先度付きキューを実現するためのヒープキューアルゴリズムの実装を提供するモジュールです。

優先度付きキュー

優先度付きキューは次の操作を提供するデータ構造です1:

挿入 : 要素を優先度付きで追加する
参照 : 優先度が最大の要素を取得する
取り出し : 優先度が最大の要素を削除し、それを返す

リストに対して複数の要素を append で追加したのちに pop で取り出すと、追加した順番と逆順で取り出されます。一方、優先度付きキューの取り出しでは追加した順番に関係なく優先度が高い順で取り出されます。

heapq モジュールは「値が小さい = 優先度が高い」とみなす優先度付きキューを実現します。
以下に具体例を示します:

from heapq import heapify, heappop, heappush

a = [4, 6, 1, 8] # リストの定義
heapify(a) # リストの中身がヒープになるように並べ替える
print(a) # [1, 6, 4, 8]
print(a[0]) # 1 : 最小値の参照

print(heappop(a))  # 最小値の取り出し
print(a) # [4, 6, 8]

heappush(a, 2) # 新たな要素の追加
print(a) # [2, 4, 8, 6]

for _ in range(len(a)):
    print(heappop(a))  # 2 -> 4 -> 6 -> 8

heappop(a) # エラーとなる!(IndexError: index out of range)
# a は空リストになっているため

上の例では以下のような処理が行われています:

  1. リスト a を定義します
  2. heapify 関数により a の中身を「ヒープ」になるように並べ替えます。この関数はあくまで並べ替えるだけであり、a は相変わらずリストであることには変わりません。ヒープの詳細は省略しますが、リストの先頭に最小値が来ていることが非常に重要です。
  3. heappop(a) によって最小値 1 が削除されます。削除後のリストもヒープ条件が満たされていることを確認してください([4, 6, 8] と先頭に最小値が来ています!)。
  4. heappush(a, 2) によって新たな要素 2 を追加します。追加後のリストは [2, 4, 8, 6] であり、やはりヒープ条件が満たされています。リストの先頭以外の要素の順番は入れ替わる可能性があることに注意してください。
  5. 4回の heappop(a) の呼び出しによって値が小さい方から取り出されます
  6. 空のリストに対する heappop はエラーになるため注意してください

計算量は以下の通りです。(N は操作対象のリストの要素数)

メソッド 計算量
heapify O(N)
heappush O(\log N)
heappop O(\log N)

💡 注意

優先度付きキューとして扱っているリストに対して appendpop などの操作を行うとリストがヒープではなくなってしまい、最小値の取得や削除が正しく行えなくなるため注意してください。

from heapq import heapify, heappop, heappush

a = [10, 20, 30, 40]
heapify(a)
a.append(0)
print(heappop(a))  # 10 (0 にはならない!)

最大値を優先する優先度付きキュー

最小値ではなく最大値を優先して取り出したい場合、多少の工夫が必要です。
以下の例のように、事前に全ての値を-1倍したリストをヒープ化すると、最小値として「-1 倍する前の最大値を -1 倍したもの」を得ることができます。したがって、得られた値をもう一度 -1 倍することで本来の最大値を得ることができます。

from heapq import heapify, heappop, heappush

a = [4, 6, 1, 8]
a = [-v for v in a] # 全ての値を -1 倍したリスト
heapify(a)

x = - heappop(a) # (-1 倍する前の) 最大値を取得し、-1 倍する
print(x) # 8
heappush(a, -2) # 2 の追加 ... 追加時にも -1 倍することに注意
y = - heappop(a)
print(y) # 6
z = - heappop(a)
print(z) # 4

最小値に紐づく要素を取得する方法

要素と値のペアが与えられている設定において、最小値の値そのものではなく、最小値を与える要素を取り出したい場合は「優先度とキーのタプルからなるリスト」をヒープ化すれば良いです。

実装例
from heapq import heapify, heappop, heappush

a = [(30, "task_A"), (10, "task_B"), (20, "task_C")]
heapify(a)

print(heappop(a)) # (10, "task_B")

heappush(a, (25, "task_D"))
print(heappop(a))  # (20, "task_C")
print(heappop(a))  # (25, "task_D")

上の例では、各タスクに値が紐づいています。(値, タスク) のタプルからなるリストをヒープ化することで、値の小さい順にタスクを取り出すことができています。
※ タプルの順番に気を付けてください。(タスク, 値) の順するとタスクが文字列として小さいものから取り出されてしまいます。

優先度付きキューを実現する仕組みや、その他のメソッドについては公式ドキュメントを参照してください。

bisect

bisectソート済みのリストに対し、指定した要素をソートされた状態を保ったまま挿入するための位置を高速に計算する機能を提供するモジュールです2

挿入位置の計算

競技プログラミングでは次の2つの関数を使いこなすことが重要になります:

  • bisect.bisect_left(v, l) : リスト3 l に値 v をソートされたまま挿入できる位置インデックスのうち、最も左のものを返します
  • bisect.bisect_right(v, l) : リスト l に値 v をソートされたまま挿入できる位置インデックスのうち、最も右のものを返します

両者は挿入する要素 v と等しい要素がリスト中に存在しなければ同じ値を返します。
※ これらの関数は実際の挿入操作をするわけではありません

利用例
from bisect import bisect_left, bisect_right

a = [10, 20, 40, 60, 90] # ソート済みのリスト

print(bisect_left(a, 30))  # 2
print(bisect_right(a, 30))  # 2

print(bisect_left(a, 60))  # 3
print(bisect_right(a, 60))  # 4

print(bisect_left(a, 99))  # 5
print(bisect_right(a, 99))  # 5

print(bisect_left(a, 5))  # 0
print(bisect_right(a, 5))  # 0
  • 30 を挿入する場合、0-indexed で 2 の位置に挿入することができます ([10, 20, 30, 40, 60, 90])。そのため、値として 2 が返ってきます
  • 60 は既にリストに存在するため、挿入箇所の候補が2つあります([10, 20, 40, 60, 60, 90])。bisect_left ではそのうちの最小値である 3 が、bisect_right では最大値である 4 が返ります。
  • 99 (リストのどの要素よりも大きい) を挿入できるのはリストの末尾です。そのためリストの要素数 = 5 が返ります。
  • 0 (リストのどの要素よりも小さい) を挿入できるのはリストの先頭です。そのため 0 が返ります。

計算量は以下の通りです。(N は操作対象のリストの要素数)

メソッド 計算量
bisect_left O(\log N)
bisect_right O(\log N)

なお、 bisect_right と全く同じ機能のメソッドが bisect という名前でも提供されていますが、_right をつけた方が取り違いを防ぎやすくなります。

参考:公式ドキュメント

活用方法

以上の説明において、「ソート済みリストへの挿入位置」を知る目的が分からないと思われるかもしれません。
実は bisect_left bisect_right の返り値は次のようにも解釈することができます:

  • bisect_left(v, l) : リスト l に含まれる v 未満の値を持つ要素の個数
  • bisect_right(v, l) : リスト l に含まれる v 以下の値を持つ要素の個数

この性質はたびたび競技プログラミングで利用することができます。すぐには納得できないかもしれませんが、前項の例と照らし合わせて成り立っていることを確認してみてください。

さらに、次のような解釈も可能です:

  • bisect_left(v, l) : リスト l に含まれる v 以上の値を持つ要素のうち最も左にあるもののインデックス (存在しない場合は len(l))
  • bisect_right(v, l) : リスト l に含まれる v より大きい値を持つ要素のうち最も左にあるもののインデックス (存在しない場合は len(l))

この性質は「挿入位置を探索する」という当初の役割をうまく言い換えたものになります。この性質を使うとソート済みリストから値を検索することが可能になります。

更に、以下のようなインデックスも検索可能です:

  • v等しい値を持つ要素のうち、最もにあるもののインデックス
  • v等しい値を持つ要素のうち、最もにあるもののインデックス
  • v 未満の値を持つ要素のうち、最もにあるもののインデックス
  • v 以下の値を持つ要素のうち、最もにあるもののインデックス

練習問題

ABC253 C - Max - Min Query
ABC217 C - Min Difference

問題

本節にはEX問題はありません。本文中にある練習問題を解いてください。

前のページ | 次のページ


  1. 優先度付きキューのアルゴリズムによっては「指定された要素の優先度を変更する」という操作もサポートしますが、heapq では提供されていません。
    (以下気になる方向け) そのような操作が必要となった場合、「キーと今の優先度の対応を表す辞書を別に保持しておく。優先度変更時には、辞書を更新し、優先度付きキュー内の元の要素に対しては何も行わず、単に新たな要素を優先度付きキューに追加する」とするのが単純な解決方法の1つです。キューから要素を取り出した際には、辞書と照らして今の優先度と一致しているかを確認し、一致していなければ無視する(再び別の要素をキューから取り出す)ことで、古い情報を参照してしまうことを防げます。 

  2. 二分探索 (binary search) という方法を使用しているためこのようなモジュール名になっています。 

  3. 本節で現れる「リスト」は全てソート済みとします。