1. ホーム
  2. パイソン

[解決済み] カスタム比較述語によるヒープク

2022-03-03 18:28:12

質問

カスタムソート述語でヒープを構築しようとしています。そこに入る値は「ユーザー定義」タイプなので、私はそれらの組み込み比較述語を変更することはできません。

というような方法はないでしょうか。

h = heapq.heapify([...], key=my_lt_pred)
h = heapq.heappush(h, key=my_lt_pred)

もっといいのは、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]

(余分な self.index の部分は、評価されたキー値がドローで、保存された値が直接比較できない場合の衝突を避けるためのものです - そうでなければ、heapq は TypeError で失敗します)