pythonのheapqが気に入らないので自分で実装しなおす(優先度付きキュー)

優先度付きキュー

って何?

って方はこちらをご参照。

ufcpp.net

こちらの本でお勉強もよろし(わかりやすかった)

アルゴリズム図鑑 絵で見てわかる26のアルゴリズム

アルゴリズム図鑑 絵で見てわかる26のアルゴリズム

ざっくり言うと、あるデータの塊から、

最小(あるいは最大)の値を取り出すために特化したデータ構造を、

ヒープ木とか二分ヒープとか優先度付き待ち行列とか優先度付きキューとかと言う。

ここらへんの表記ゆれはよくわからない。

またデータを追加するときにenqueue(エンキュー)、

最小値を取り出す(そして削除)のをdequeue(デキュー)という。

特徴として、先頭の値を取り出すのがO(1)と高速だし、 追加や削除もO(log n)とそこそこ高速である。

また最後までdequeueし続けるとヒープソートの完成である。

pythonの優先度付きキュー

pythonにはheapqというライブラリが備わっており、

enqueueはheappush,dequeueはheappopで備わっているのだが、

下記事項が気に食わないので作り直した。

  • オブジェクト志向ではない

    事前に配列を自分で宣言し、heappushやheappopに配列を引数で突っ込む必要があり、直感的でない

  • 降順をサポートしない

    優先度は昇順しか使えないので、降順にしたい場合は別途正負反転させる等の工夫が必要

  • heappush,heappopってなんやねん

    enqueue,dequeueのほうがかっこよくない?

ということで実装。

実装

pure pythonで実装。

class heap:
    def __init__(self,order="asc"):
        self._heap = []
        self.order=order
    def get_head(self):
        return self._heap[0] if self.order == "asc" else -self._heap[0]
    def get_heap(self):
        return self._heap if self.order == "asc" else list(map(lambda x:-x,self._heap))
    def enqueue(self,items):
        if type(items) != list:
            items = [items]
        # 降順ソートの場合は正負反転
        if self.order!="asc":
            items = list(map(lambda x:-x,items))
        for item in items:
            self._heap.append(item)
            self_idx = len(self._heap)-1
            parent_idx = self_idx//2
            while self_idx > 0:
                if self._heap[self_idx] <= self._heap[parent_idx]:
                    self._heap[self_idx], self._heap[parent_idx] = self._heap[parent_idx], self._heap[self_idx]
                    self_idx = parent_idx
                    parent_idx = self_idx//2
                else:
                    break
    def dequeue(self):
        head = self._heap[0]
        self._heap[0] = self._heap[-1]
        self._heap.pop(-1)
        self_idx = 0
        child_idx = [1,2]
        while True:
            # 子供が存在しなかったらbreak
            if child_idx[0] > len(self._heap)-1:
                break
            # 左の子供しか存在しない場合
            elif child_idx[0] == len(self._heap)-1:
                # 左の子供のほうが大きい場合
                if self._heap[self_idx] < self._heap[child_idx[0]]:
                    break
                else:
                    self._heap[self_idx],self._heap[child_idx[0]] = self._heap[child_idx[0]],self._heap[self_idx]
                    self_idx = (self_idx*2)+1
                    child_idx = [self_idx*2+1,self_idx*2+2]
            # 両方の子供が存在する場合
            else:
                # 子供のほうが大きい場合
                if self._heap[self_idx] < min([self._heap[child_idx[0]],self._heap[child_idx[1]]]):
                    break
                # 子供のほうが小さい場合
                else:
                    # 左の子供のほうが小さい場合は左の子供と親の値を交換
                    if self._heap[child_idx[0]] <= self._heap[child_idx[1]]:
                        self._heap[self_idx],self._heap[child_idx[0]] = self._heap[child_idx[0]],self._heap[self_idx]
                        self_idx = (self_idx*2)+1
                        child_idx = [self_idx*2+1,self_idx*2+2]
                    # 右の子供のほうが小さい場合は右の子供と親の値を交換
                    else:
                        self._heap[self_idx],self._heap[child_idx[1]] = self._heap[child_idx[1]],self._heap[self_idx]
                        self_idx = (self_idx*2)+2
                        child_idx = [self_idx*2+1,self_idx*2+2]
        return head  if self.order == "asc" else -head

使い方

# heapインスタンス生成
heap = heap()

# [10,9,8,7,6,5,4,3,2,1]というデータを突っ込む
heap.enqueue([10-i for i in range(10)])

# 先頭のデータを取得(参照のみでデータ変化なし)
heap.get_head() # 1

# heap全体の取得(参照のみでデータ変化なし)
heap.get_heap() # [1, 2, 3, 5, 4, 8, 9, 6, 10, 7] ←先頭が最小値であることと、木構造で見たときに「親の値<子供の値」のみ保証

# データを突っ込むその2(先頭に0を突っ込む)
heap.enqueue(0)

# 先頭が書き換わったことの確認
heap.get_head() # 0

# dequeue
print(heap.dequeue()) # 0が返り0は消滅

# 先頭が書き換わる
heap.get_head() # 1

# heapsort
heap_sort_result = []
for i in range(len(heap.get_heap())):
    heap_sort_result.append(heap.dequeue())
print(heap_sort_result) # [1, 2, 3, 4, 5, 6, 7, 8, 9, 10]

こんな感じで使えます。

同等のことをheapqでやってみる。

import heapq
# heapインスタンス生成相当(ただのリスト宣言)
Q = []
# データを突っ込む
for i in [10-i for i in range(10)]:
    heapq.heappush(Q,i)

# 先頭の値を取得
print(Q[0])

# heap全体の取得
print(Q)

# データを突っ込むその2(先頭に0を突っ込む)
heapq.heappush(Q,0)

# 先頭が書き換わったことの確認
print(Q[0])

# dequeue
heapq.heappop(Q)

# heapsort
heap_sort_result = []
for _ in range(len(Q)):
    heap_sort_result.append(heapq.heappop(Q))
print(heap_sort_result)

毎回リストを引数につっこまなくてよくなりました。

また、降順はこんな感じ。

# インスタンス生成時に降順を指定
heap = heap(order = "desc")

# 昇順にデータを突っ込む
heap.enqueue([i+1 for i in range(10)]) # [1,2,3,4,5,6,7,8,9,10]を格納

#10が先頭に来ている
heap.get_head() 

# heap sort(降順)
heap_sort_result = []
for i in range(len(heap.get_heap())):
    heap_sort_result.append(heap.dequeue())
print(heap_sort_result) # [10, 9, 8, 7, 6, 5, 4, 3, 2, 1]

これをheapqでやるとこうなる

import heapq

# インスタンス生成時に降順を指定(という名のただのリスト生成)
Q = []

# 昇順にデータを突っ込む
for i in [i+1 for i in range(10)]:
    heapq.heappush(Q,-i)

#10が先頭に来ている
print(-Q[0])

# heap sort(降順)
heap_sort_result = []
for i in range(len(Q)):
    heap_sort_result.append(heapq.heappop(Q))
heap_sort_result = list(map(lambda x:-x,heap_sort_result))
print(heap_sort_result) # [10, 9, 8, 7, 6, 5, 4, 3, 2, 1]

heapqは降順ソートをサポートしていないので、 毎回正負反転が必要で頭がこんがらがってしまう。

なので、こちらをオススメをしたい。

処理速度

とはいえ、処理速度はどうだろう。 1~1,000,000を突っ込んでheap sortしてみる

今回作ったやつ

%%time
heap = heap(order = "desc")
heap.enqueue([i+1 for i in range(1000000)])
heap_sort_result = []
for i in range(len(heap.get_heap())):
    heap_sort_result.append(heap.dequeue())

Wall time: 31.2 s

heapq

%%time
import heapq
# インスタンス生成時に降順を指定(という名のただのリスト生成)
Q = []
# 昇順にデータを突っ込む
for i in [i+1 for i in range(1000000)]:
    heapq.heappush(Q,-i)
# heap sort(降順)
heap_sort_result = []
for i in range(len(Q)):
    heap_sort_result.append(heapq.heappop(Q))
heap_sort_result = list(map(lambda x:-x,heap_sort_result))

Wall time: 1.16 s

ぐえ、ボロ負け。

仕方ないのでnumbaで高速化を図る…

のは次回以降。