📌

pythonでカスタムクラスを配列に入れてmin/max/sortをする

2022/01/26に公開

Python

比較

class

tech

例えば以下のカスタムクラスA。v1とv2の値を持つ。

class A:
    def __init__(self, v1, v2):
        self.v1 = v1
        self.v2 = v2

    def __str__(self):
        return f'[v1: {self.v1}, v2: {self.v2}]'

このクラスのインスタンスを幾つか配列に入れる。

a = [A(1, 10), A(5, 5), A(9, 2)]

この配列からv1が最小、最大のインスタンスを取り出したい。

単純にmin/max関数に配列を渡してもエラーになる。

>>> print(min(a))
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
TypeError: '<' not supported between instances of 'A' and 'A'

原因はクラスA同士の比較（<）がサポートされていないため。

比較演算子の特殊メソッドをオーバーライド

<は特殊メソッドの__lt__をオーバライドすることでサポートできる。

Aで__lt__をサポートしてみる。

class A:
    def __init__(self, v1, v2):
        self.v1 = v1
        self.v2 = v2

    def __str__(self):
        return f'[v1: {self.v1}, v2: {self.v2}]'

    def __lt__(self, other):
        return self.v1 < other.v1

確認してみる。

>>> print(min(a))
[v1: 1, v2: 10]

できた。

maxの場合も大丈夫。

>>> print(max(a))
[v1: 9, v2: 2]

sortもできてる。

>>> a.sort()
>>> print(a[0], a[1], a[2])
[v1: 1, v2: 10] [v1: 5, v2: 5] [v1: 9, v2: 2]

しかし、これだとAの他のプロパティv2で比較したい場合に困る。

keyを指定して比較する

そんな場合にはmin/max/sortでkeyを指定して比較を行う。

>>> print(min(a, key=lambda x: x.v2))
[v1: 9, v2: 2]
>>> print(max(a, key=lambda x: x.v2))
[v1: 1, v2: 10]

sortの場合も同様にkeyを指定することが可能。

>>> a.sort(key=lambda x: x.v2)
>>> print(a[0], a[1], a[2])
[v1: 9, v2: 2] [v1: 5, v2: 5] [v1: 1, v2: 10]

これならカスタムクラスのどのプロパティでも比較ができる。

まとめ

カスタムクラスを作った場合、比較などに使用する値が1つに決まっているのであれば比較演算子の特殊メソッドをオーバーライドで対応しても良い。しかし、インスタンスを様々なプロパティで比較したい場合にはmin/max/sortなどの比較する関数にkeyを渡して比較するのが良さそう。

比較演算子の特殊メソッドをオーバーライド

keyを指定して比較する

まとめ

Discussion