カスタム比較述語を使用した heapq 質問する

によるheapq ドキュメントヒープ順序をカスタマイズする方法は、ヒープ上の各要素をタプルにして、最初のタプル要素を通常の Python 比較を受け入れる要素にすることです。

heapq モジュールの関数は少々扱いにくく (オブジェクト指向ではないため)、常にヒープ オブジェクト (ヒープ化されたリスト) を最初のパラメータとして明示的に渡す必要があります。関数を指定してヒープをオブジェクトkeyとして提示できる非常にシンプルなラッパー クラスを作成することで、一石二鳥を実現できます。

以下のクラスは内部リストを保持します。各要素はタプルで、その最初のメンバーはキーであり、keyヒープインスタンス化時に渡されるパラメータを使用して要素の挿入時に計算されます。

# -*- coding: utf-8 -*-
import heapq

class MyHeap(object):
    def __init__(self, initial=None, key=lambda x:x):
        self.key = key
        self.index = 0
        if initial:
            self._data = [(key(item), i, item) for i, item in enumerate(initial)]
            self.index = len(self._data)
            heapq.heapify(self._data)
        else:
            self._data = []

    def push(self, item):
        heapq.heappush(self._data, (self.key(item), self.index, item))
        self.index += 1

    def pop(self):
        return heapq.heappop(self._data)[2]