ごりちゃんがゆく

競プロをするゴリラの精進記録

Pythonでheapqから大きい順に取り出したいときにもバグらせにくいやつを書いた

Python スニペット

表題の通りです。

お気持ちを説明すると、ご存知 heapq は優先度付きキューと呼ばれていて、適当に要素を追加したり取り出したりしても常に取り出す値は最小を保っている便利なやつです。 しかも追加・取り出しはともにO(logN)でとても実用的！

しかし欠点があって、これの逆バージョン、つまり常に最大値を取り出したいときには対応していないんです。

一応それを実現するテクはあって、全部符号を反転してheapqに突っ込めばデカい方から順に取り出せます！
でも追加したり取り出したりするたびに符号を反転するの結構メンドイし、バグらせやすい。
コードは運よくバグらなくても、実装中に頭がバグる。

そんな悩みを解決するために、classでラップしちゃおう！と思い立ちました。

大したことないコードですが、あると地味に嬉しいだろうと思ったので、書いてみました。

import heapq
class Heapq:
    def __init__(self, arr, desc=False):
        if desc:
            arr = [-a for a in arr]
        self.sign = -1 if desc else 1
        self.hq = arr
        heapq.heapify(self.hq)
 
    def pop(self):
        return heapq.heappop(self.hq) * self.sign
 
    def push(self, a):
        heapq.heappush(self.hq, a * self.sign)
 
    def top(self):
        return self.hq[0] * self.sign

使い方

まず上のコードをコピペする。ドーン。

初期化

q = Heapq(arr, desc) のように書けばいいです

  • 第1引数 arr は初期化に使う配列です 空っぽから始める場合は []でいいです
  • 第2引数 desc は大きい順に取り出すなら True、小さい順ならFalse です。Falseは省略可です。

pop()

q.pop() のようにすると値が返ります
一番小さいの or 一番大きいのを集合から取り出します。
初期化時 desc=Trueを指定していたなら大きい方が出てきます。
取り出された要素は集合から消えます。

push()

q.push(a) のようにすると集合に aを追加します

top()

q.top() のようにすると値が返ります
一番小さいの or 一番大きいのを参照できます
初期化時 desc=Trueを指定していたなら大きい方が出てきます。
pop()に似てますが、参照するだけで、集合から値が消えることはありません。
O(1)でめっちゃ速いです

個人的に気に入っている点

元の heapq.heappush(hq, 3) みたいな書き方くどくない？
hq.push(3) みたいに書けるので楽ちん！

試し切り

heapq と逆向きのheapqとを2個持つ問題だったのでちょうど良かった。
atcoder.jp

注意点

本家のheapqのように配列やタプルの配列には対応していなくて、単純な数値の配列のみに対応しています
もし改良してくれる人がいたら教えてください

→→ここから下は2019/6/25追記

情報提供1

@tatyam_prime さんが改良版を考えてくださいました

上のリンク先から引用します

import heapq
class Reverse:
    def __init__(self, val):
        self.val = val
        
    def __lt__(self, other):
        return self.val > other.val
        
    def __repr__(self):
        return repr(self.val)
    
class Heapq:
    def __init__(self, arr, desc = False):
        if desc:
            for i in range(len(arr)):
                arr[i] = Reverse(arr[i])
        self.desc = desc
        self.hq = arr
        heapq.heapify(self.hq)
 
    def pop(self):
        if self.desc: return heapq.heappop(self.hq).val
        else: return heapq.heappop(self.hq)
 
    def push(self, a):
        if self.desc: heapq.heappush(self.hq, Reverse(a))
        else: heapq.heappush(self.hq, a)
 
    def top(self):
        if self.desc: self.hq[0].val
        else: return self.hq[0]

見た感じ、desc=Trueのときは Reverse() オブジェクト同士を(おそらくheapq内部で)比較するようにしていて、Reverse()オブジェクトは __lt__() がオーバーライドされてて本来と逆向きに比較されるんでしょうかね？
かなりテクいです
これで数値の配列に限らず、文字列やタプルや配列の配列にも使えて、汎用的ですね！

情報提供2

@juppyjappy さんからさらに情報提供がありました！

なんとheapq には最大値から取り出すバージョンが存在しない、というのは思い込みで、一応できる方法があるらしい！！ 完全に知らなかった！

もう少し書くと、heapify() に当たるものは heapify_max()という名前で存在している！
だけどheappop() や heappush() に当たるものは存在しない笑！
惜しい部品はあるからそれでやってね、みたいな感じらしい
上のブログ記事から引用します

import heapq
#これを加える！！！！
def _heappush_max(heap, item):
    heap.append(item)
    heapq._siftdown_max(heap, 0, len(heap)-1)
def _heappop_max(heap):
    """Maxheap version of a heappop."""
    lastelt = heap.pop()    # raises appropriate IndexError if heap is empty
    if heap:
        returnitem = heap[0]
        heap[0] = lastelt
        heapq._siftup_max(heap, 0)
        return returnitem
    return lastelt

こうやってやると、heapify_max() で構築した降順heapに対して、heappop() と heappush() ができるらしい！
ほえ～、これそのまま使えるように組み込んでくれ～～
元記事はheapについて詳しく書いてある記事なので、気になった方は読んでみてください！！

結局どれを使えばいいの？

好きなのを使いましょう！！