24 November 2013

python の sorted 関数のインタフェース

昔の python の sorted 関数はソートしたいリストと比較関数を渡すインタフェースだったが、どうも python3 では閉じられたようだ。代わりに key と reverse という引数を渡す。key は sorted が比較を行う前のプレフィルタで reverse が asc/desc のオーダー順を切り替えるスイッチだ。

たとえば文字列を case-insensitive で辞書順にソートしたい場合、key に str.lower を渡すと、lower したあとの文字列で比較してくれる。

>>> sorted("This is a test string from Andrew".split(), key=str.lower)
['a', 'Andrew', 'from', 'is', 'string', 'test', 'This']

オブジェクトを特定のキーでソートしたい場合はこんなかんじだ。

>>> student_tuples = [
...         ('john', 'A', 15),
...         ('jane', 'B', 12),
...         ('dave', 'B', 10),
... ]
>>> sorted(student_tuples, key=lambda student: student[2])   # sort by age
[('dave', 'B', 10), ('jane', 'B', 12), ('john', 'A', 15)]

このようなよく使われるケースのためのショートカットもある。itemgetter, attrgetter に比較したいインデックスやキーを指定する。複数指定もできる。

>>> from operator import itemgetter, attrgetter

>>> sorted(student_tuples, key=itemgetter(2))
[('dave', 'B', 10), ('jane', 'B', 12), ('john', 'A', 15)]

>>> sorted(student_objects, key=attrgetter('age'))
[('dave', 'B', 10), ('jane', 'B', 12), ('john', 'A', 15)]

>>> sorted(student_tuples, key=itemgetter(1,2))
[('john', 'A', 15), ('dave', 'B', 10), ('jane', 'B', 12)]

>>> sorted(student_objects, key=attrgetter('grade', 'age'))
[('john', 'A', 15), ('dave', 'B', 10), ('jane', 'B', 12)]

何も考えずに古いスタイルの cmp 関数を key で使えるように変換するフィルタも提供されている。functools モジュールの cmptokeys という関数だ。実装はこんなかんじで、比較演算子を cmp 関数でオーバーライドするという内容だ。ただ毎回クラスを作って返すのでパフォーマンス的には不利そうだ。

def cmp_to_key(mycmp):
    'Convert a cmp= function into a key= function'
    class K(object):
        def __init__(self, obj, *args):
            self.obj = obj
        def __lt__(self, other):
            return mycmp(self.obj, other.obj) < 0
        def __gt__(self, other):
            return mycmp(self.obj, other.obj) > 0
        def __eq__(self, other):
            return mycmp(self.obj, other.obj) == 0
        def __le__(self, other):
            return mycmp(self.obj, other.obj) <= 0
        def __ge__(self, other):
            return mycmp(self.obj, other.obj) >= 0
        def __ne__(self, other):
            return mycmp(self.obj, other.obj) != 0
    return K

この key, reverse のインタフェースは python2.4 の頃からあるらしく、けっこう古い。

参考