pythonのheapqが気に入らないので自分で実装しなおす(優先度付きキュー)
優先度付きキュー
って何?
って方はこちらをご参照。
こちらの本でお勉強もよろし(わかりやすかった)
- 作者: 石田保輝,宮崎修一
- 出版社/メーカー: 翔泳社
- 発売日: 2017/06/06
- メディア: 単行本(ソフトカバー)
- この商品を含むブログを見る
ざっくり言うと、あるデータの塊から、
最小(あるいは最大)の値を取り出すために特化したデータ構造を、
ヒープ木とか二分ヒープとか優先度付き待ち行列とか優先度付きキューとかと言う。
ここらへんの表記ゆれはよくわからない。
またデータを追加するときに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で高速化を図る…
のは次回以降。