カスタムソート関数の作成方法(AtCoder ABC308)

本記事の内容

配列の要素を比較する際に要素での比較だけではなく処理を行ったうえで比較する方法を紹介します

競技プログラミングサイトであるAtCoderのABC308 C問題を使用して解説をしていきます。

問題文

この問題を解くにあたって、単純に期待値を求めた後にソートしたいなという発想がまず最初に思い浮かぶと思います。

愚直に書いたコード(正しく動かない)

n = int(input())
ab_n = [list(map(int, input().split())) for _ in range(n)]

arr = []
for i in range(n):
    a, b = ab_n[i]
    expected_val = a / (a + b)
    arr.append((i, expected_val))

sorted_arr = sorted(arr, key=lambda x: x[1], reverse=True)

ans = [i + 1 for i, _ in sorted_arr]
print(*ans)

解決方法

上記のコードが正しく動かない理由として、浮動小数点が登場しており、ここで誤差が発生しています。
これを解決する方法として、分母を払い通分をして整数での比較をすることで誤差が発生することなく求めることが出来ます。

(例)

\displaylines{
\frac{99}{100} と \frac{100}{101} 
}

を比較する際には分母を払い差を求めることで正しく値を比較出来ます。

分母を払う

99 \times 101 - 100 \times 100 
= 9999-10000
=-1

今回であれば

\frac{100}{101}

のほうが大きいことがわかります。

これをコードに落とし込むために要素ごとの比較をする必要があります。そのためにカスタムソート関数を作成します。

カスタムソートのコード

def key_func(i, j):
    """
    条件に基づいて2つのインデックスを比較

    2つのリストAとBの要素に基づいて値を計算し比較結果に基づいて-1、1を返す。
    -1のときはそのまま、1のときはiとjを入れ替える。
    """
    diff = A[i] * (A[j] + B[j]) - A[j] * (A[i] + B[i])
    if diff > 0:
        return -1
    elif diff < 0:
        return 1
    else:
        return -1 if i < j else 1

-1,1についてですが、後に登場するソート関数の内部の処理で-1であれば入れ替えない、1であれば入れ替えるという処理が行われています。

正解したコード

from functools import cmp_to_key


def key_func(i, j):
    """
    条件に基づいて2つのインデックスを比較

    2つのリストAとBの要素に基づいて値を計算し比較結果に基づいて-1、1を返す。
    -1のときはそのまま、1のときはiとjを入れ替える。
    """
    diff = A[i] * (A[j] + B[j]) - A[j] * (A[i] + B[i])
    if diff > 0:
        return -1
    elif diff < 0:
        return 1
    else:
        return -1 if i < j else 1


n = int(input())
A, B = [0] * n, [0] * n
for i in range(n):
    A[i], B[i] = map(int, input().split())

ans = list(range(n))
ans.sort(key=cmp_to_key(key_func))

a = [i + 1 for i in ans]
print(*a)

今回はインデックスの順番を求める必要があったためA,Bを定数として定義して求めました。

さいごに

ソート関数の関数は引数に関数を渡すので慣れないとなかなかすぐに書けないので沢山書いていきたいと思います。

(追記)
久しぶりにTEXを書いたのですが、QiitaはTEXにも対応していて本当に便利だなと実感しました。