NAKED WETWARE

「わがままになるのが怖い奴に宇宙は拓けねェさ」

Pythonで独自比較関数を用いてソートに適用する

Python C#

Pythonでリストをソートするとき、独自の比較関数を使いたかった。

そのことについて調べたことをまとめます。

TL;DR

  • Python2.xまではsort関数のcmp キーワード引数に独自compare関数を指定すればよかった
  • Python3.xでは、sort関数からcmp キーワード引数が完全に削除された
  • Python3.xで独自compare関数custom_compareを使う場合はsorted(your_list, key=functools.cmp_to_key(custom_comare))とする
  • 単純な二重配列の場合 sorted(your_list, key=operator.itemgetter(0,1,2))のような書き方もできる

事の発端

ソースコードを解析するスクリプトを書いているときに、以下のような二重配列のデータ構造を作った。

list_buff = [
    ['fuga', 69, 82],
    ['fuga', 85, 92],
    ['hoge', 93, 99],
    ['fuga', 115, 137],
    ['hoge', 140, 152]
]

これを、

  • list_buff[n][0]は辞書順
  • list_buff[n][1]は昇順

といった具合にソートしたかった。

具体的には、以下のように直したい。

list_buff = [
    # 辞書順,  昇順,  スコープ外
    ['fuga',    69,    82],
    ['fuga',    85,    92],
    ['fuga',   115,   137],
    ['hoge',    93,    99],
    ['hoge',   140,   152]
]

C#とかだとcompareクラスを作れば、わりかしsortの自由が利く。

以下はC#のイメージ。

using System;
using System.Collections;
using System.Collections.Generic;

namespace CustomSort
{
    public class RowData
    {
        public string Label { get; set; }
        public int Start { get; set; }
        public int End { get; set; }

        public static RowData Create(string label, int start, int end)
        {
            return new RowData()
            {
                Label = label,
                Start = start,
                End   = end
            };
        }
    }

    public class CustomCompare : System.Collections.Generic.IComparer<RowData>
    {
        public int Compare(RowData elm1, RowData elm2)
        {
            int diff = elm1.Label.CompareTo(elm2.Label);
            if (diff != 0) { return diff; }

            return elm1.Start.CompareTo(elm2.Start);
        }
    }

    class EntryPoint
    {
        static void Main(string[] args)
        {
            // メイン処理イメージ
            List<RowData> listBuff = new List<RowData>()
            {
                RowData.Create("fuga", 69, 82),
                RowData.Create("fuga", 85, 92),
                RowData.Create("hoge", 93, 99),
                RowData.Create("fuga", 115, 137),
                RowData.Create("hoge", 140, 152)
            };
            
            CustomCompare comp = new CustomCompare();
            listBuff.Sort(comp);

            foreach(RowData rd in listBuff)
            {
                System.Console.WriteLine(@"{0}: {1}: {2}", rd.Label, rd.Start, rd.End);
            }
        }
    }
}

Pythonでも、以下のようなことをやりたい。

if __name__ == '__main__':
    list_buff = get_datas()

    def custom_compare(elm1, elm2):
        if elm1[0] > elm2[0]:
            return 1
        elif elm1[0] < elm2[0]:
            return -1
        else:
            return elm1[1] - elm2[1]

    # # こういうイメージ
    # list_buff.sort(compare=custom_compare)

Pythonにもないかしらと調べた。

結果

公式ドキュメントを調べて sort(cmp=custom_function) で実現できそうだとわかった。

しかし、上記の文法は2.xまでしかサポートされておらず、3.xではcmp キーワード引数は完全に削除されたとのこと。

3.xで、独自に拡張したcompare関数を使いたい場合、keyキーワード引数にcompare関数をラップしたクラスを渡せばいい。

if __name__ == '__main__':
    list_buff = [
        ['fuga', 69, 82],
        ['fuga', 85, 92],
        ['hoge', 93, 99],
        ['fuga', 115, 137],
        ['hoge', 140, 152]
    ]

    def custom_compare(elm1, elm2):
        if elm1[0] > elm2[0]:
            return 1
        elif elm1[0] < elm2[0]:
            return -1
        else:
            return elm1[1] - elm2[1]

    def cmp2key(comp_func):
        class CustomCompare:
            def __init__(self, obj, *arg):
                self.obj = obj
            # 以下の__eq__のように、eq, ne, lt, gt, le, ge を定義する
            def __eq__(self, other):
                return comp_func(self.obj, other.obj) == 0
            ...
        # cmp2keyのreturn
        return CustomCompare

    # 独自のソートを実施
    list_buff.sort(key=cmp2key(custom_compare))

ミソは、C#の時と似たような感じ。compare関数を持つクラスを作成すること。

ただ、C#と違うのは、拡張比較演算子をオーバーライドすること。

正直、めんどくさい。

代わりに, functoolsモジュールにあるcmp_to_key関数を使えば、実現したいことを実現できることがわかった。1

from functools import cmp_to_key



if __name__ == '__main__':
    list_buff = [
        ['fuga', 69, 82],
        ['fuga', 85, 92],
        ['hoge', 93, 99],
        ['fuga', 115, 137],
        ['hoge', 140, 152]
    ]

    def custom_compare(elm1, elm2):
        if elm1[0] > elm2[0]:
            return 1
        elif elm1[0] < elm2[0]:
            return -1
        else:
            return elm1[1] - elm2[1]

    # 独自のソートを実施
    list_buff.sort(key=cmp_to_key(custom_compare))

簡単。

ただ、cmp_to_keyが書かれた箇所の少し上に、もっと単純な書き方が載っていた2。

特に今回のような二重配列の場合、以下の使い方が最も簡単。

from operator import itemgetter

if __init__ == '__main__':
    list_buff = [
        ['fuga', 69, 82],
        ['fuga', 85, 92],
        ['hoge', 93, 99],
        ['fuga', 115, 137],
        ['hoge', 140, 152]
    ]

    list_buff.sort(key=itemgetter(0,1))

参考資料

  1. cmp パラメータを利用した古い方法↩

  2. operator モジュール関数↩