heap queueを自分で実装してみる

プライオリティキュー実装としてよくつかわれるpythonのheapqモジュールのソースは、ばらばらのリストをheap化するheapifyのようなものがあったり、細かい効率化などもあって、そのソースコードは結構な大きさです。

基本はそうたいそうなものではないので、訓練として自らの手で実装してみました。heapを空リストから追加前提にし、アルゴリズム部分をなるべくシンプルに書いてみると、このくらいになりました。

heap操作の理解のためのメモ

heapは、二分木の配列表現の一種:

  • インデックスが1始まりだとすると、インデックスnの子は2nと2n+1になる
    • たとえば、1の子は2と3、2の子は4と5、3の子は6と7のようにかぶることは無い
    • CやPythonのように0始まりの場合は、2n+1と2n+2
  • さらに、すべてのノードで、親の値は子の値より必ず小さい必要がある
    • 先頭が最小値
    • (そうなるように、push/popで調整する)

push操作は以下:

  • まず配列の最後に値を追加する
  • 最後をカーソルにしてループ
    • 親を見て、親のほうが大きければ、swapする。そうでなければ終了。
    • カーソルを移して終了まで繰り返す

pop操作は以下:

  • 先頭を結果として保存しておく
  • 最後の値をpopする
  • pop後空なら終了し、結果を返す
  • (そうでなければ)popした値を配列の先頭へ上書きする
  • 先頭をカーソルにしてループ
    • 子の値の小さいほうと比べ、それより大きければその子の値と入れ替える
    • 子がないか、子より小さければ終了し、結果を返す
    • カーソルをswapした子に移して終了まで繰り返す

pushは底から頭にswapしつつ上がる感じで、popは頭から底へswapしつつ下がっていく感じでしょうか。どちらのswap回数も最大log2(size)回です。