pythonで優先度キュー(ヒープキュー)を使う

pythonで優先度キュー(ヒープキュー)を使う

先日、AtCoder Beginner Contest 141 に参加しました。
6問中3問しか解けませんでしたが、アルゴリズムを考えながら実装するのはとても楽しかったです。
[コンテストページ]
https://atcoder.jp/contests/abc141

コンテスト終了後、すぐに回答を公開してくれるので確認したところ、問題Dの解法に「優先度付きキュー」とあったのですが勉強不足だったため調べた事を残しておきます。

優先度キューとは

優先度付きキュー(ゆうせんどつき -、: priority queue)は、以下の4つの操作をサポートする抽象データ型である。
キューに対して要素を優先度付きで追加する。
・最も高い優先度を持つ要素をキューから取り除き、それを返す。
・(オプション) 最も高い優先度を持つ要素を取り除くことなく参照する。
・(オプション) 指定した要素を取り除くことなく優先度を変更する

https://ja.wikipedia.org/wiki/%E5%84%AA%E5%85%88%E5%BA%A6%E4%BB%98%E3%81%8D%E3%82%AD%E3%83%A5%E3%83%BC

一般的なキューは、キューから取り出す際は追加した順番に依存するのですが、優先度キューは優先度の高い順に取り出せるデータ構造です。
問題Dでは「金額が最も高い商品を取り出し、半額にしてキューに戻す」という処理を繰り返し行う必要があるので、優先度キューの出番になります。

pythonにおける優先度キュー

調べてみるとpythonにも優先度キューを扱えるライブラリが組み込まれていることがわかったのですが、そのライブラリではデキュー時(キューを取り出す)優先度が最小の値を取り出すという仕様でした。
https://docs.python.org/ja/3/library/heapq.html

今回は優先度(金額)が最大の値を取得したいのでひと工夫入れないと使えません。

優先度が最大の値を取得可能にする

色々調べていると、以下の記事にpython, pypy での実装方法と仕組みの説明を見つけたのでそちらを参考に実装しました。
https://juppy.hatenablog.com/entry/2019/04/05/%E8%9F%BB%E6%9C%AC_Python_%E3%83%97%E3%83%A9%E3%82%A4%E3%82%AA%E3%83%AA%E3%83%86%E3%82%A3%E3%82%AD%E3%83%A5%E3%83%BC_2_heapq_%E7%AB%B6%E6%8A%80%E3%83%97%E3%83%AD%E3%82%B0%E3%83%A9%E3%83%9F%E3%83%B3

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

理解するためにライブラリの実装(https://github.com/python/cpython/blob/3.7/Lib/heapq.py)を確認したところ、 デキュー用の関数_heappop_max(heap)は定義されているようですが、なぜかエンキュー用の関数が見当たりませんでした。
エンキュー時にツリーを走査して、然るべき位置に値を格納する必要があるはずなのですが、なぜか中途半端な状態になっているようです。

以上、pythonで優先度キュー(ヒープキュー)を使う方法の紹介でした。