pythonで、比較できないオブジェクトを優先度を指定して優先度キューに格納する方法

自作のクラスなど、それ自体は比較可能でないオブジェクトを、クラス外で優先度を用いて優先度キューに格納したいときどうすればいいのか？方法は簡単ですがいくつかつまずいた点があったのでまとめました。

解決方法

以下のように、優先度付きキューに("優先度", "自作クラスのインスタンス")を格納することで達成できます。

import heapq

# 優先度付きキュー
priority_queue = []

# 自作のクラス
class myclass:
    def __init__(self, name):
        self.name = name

i1 = myclass('i1')
i2 = myclass('i2')
i3 = myclass('i3')

# （優先度、インスタンス）というタプルを要素として追加する
heapq.heappush(priority_queue, (1, i1))
heapq.heappush(priority_queue, (2, i2))
heapq.heappush(priority_queue, (3, i3))

# 優先度が低い順に要素を取り出す
while priority_queue:
    priority, item = heapq.heappop(priority_queue)
    print(f'Priority: {priority}, Item: {item.name}')
    
# 結果
'''
Priority: 1, Item: i1
Priority: 2, Item: i2
Priority: 3, Item: i3
'''

同じ優先度のインスタンスを格納できないエラー

しかし、上記の方法では以下のように、同じ優先度を異なるインスタンスに付与した場合にエラーが起こります。
エラーの原因は、「i1とi2の大小関係を求めるときに、タプルの第一要素をまず比べたが値が同じだったのでタプルの第二要素を比べたが、myclassはそもそも比較不可能である」からです。

import heapq

# 優先度付きキュー
priority_queue = []

# 自作のクラス
class myclass:
    def __init__(self, name):
        self.name = name

i1 = myclass('i1')
i2 = myclass('i2')
i3 = myclass('i3')

# （優先度、インスタンス）というタプルを要素として追加する
heapq.heappush(priority_queue, (1, i1))
heapq.heappush(priority_queue, (1, i2)) # さっきは(2, i2)を格納していた
heapq.heappush(priority_queue, (3, i3))

# 優先度が低い順に要素を取り出す
while priority_queue:
    priority, item = heapq.heappop(priority_queue)
    print(f'Priority: {priority}, Item: {item.name}')

# 結果
'''
Traceback (most recent call last):
  File "heapqueue.py", line 17, in <module>
    heapq.heappush(priority_queue, (1, i2))
TypeError: '<' not supported between instances of 'myclass' and 'myclass'
'''

エラーの解決方法

上記のエラーの解決方法は、クラス同士を比較可能にすることです。今回は、タプルの第一要素で与えている優先度でさえソートされていれば、同一優先度を持つmyclassインスタンス同士の順番は気にしないものとします。

以下のように、クラスmyclassに比較メソッドを追記することで解決できます。

import heapq

# 優先度付きキュー
priority_queue = []

# 自作のクラス
class myclass:
    def __init__(self, name):
        self.name = name
    def __lt__(self, other):   #比較メソッドを追記
        return True

i1 = myclass('i1')
i2 = myclass('i2')
i3 = myclass('i3')

# （優先度、インスタンス）というタプルを要素として追加する
heapq.heappush(priority_queue, (1, i1))
heapq.heappush(priority_queue, (1, i2))
heapq.heappush(priority_queue, (3, i3))

# 優先度が低い順に要素を取り出す
while priority_queue:
    priority, item = heapq.heappop(priority_queue)
    print(f'Priority: {priority}, Item: {item.name}')

# 結果
'''
Priority: 1, Item: i2
Priority: 1, Item: i1
Priority: 3, Item: i3
'''