はじめに
競プロ精進中に「計算量的には間に合いそうなのにTLEする...」と悩んでいたところ、前準備のsort部分がボトルネックになっていたと発覚。
ハマったポイントを共有します。
tupleの配列のsort
key指定あり/なしの違いについて、簡潔に説明します。
key指定あり
a = [(3,3), (2,2), (3,1), (4,2)]
a.sort(key=lambda x: x[0])
print(a) # => [(2, 2), (3, 3), (3, 1), (4, 2)]
指定したkeyの大小によって、sortを行います。keyの要素が等しければ、その順序は保持されます。
sort() メソッドは安定していることが保証されています。ソートは、等しい要素の相対順序が変更されないことが保証されていれば、安定しています。（公式ドキュメント）
sort(key=lambda x: (x[0],x[1])) のようにkeyをtupleにして複数指定することもできますが、
本記事では単一の要素をkey指定することを考えます。
key指定なし
何もkeyを指定しないと、
「tupleの0番目の要素で大小比較」→「同じなら1番目の要素で大小比較」→「...」
という方法でsortが行われます。
a = [(3,3), (2,2), (3,1), (4,2)]
a.sort()
print(a) # => [(2, 2), (3, 1), (3, 3), (4, 2)]
このとき内部では、tuple同士の比較演算が行われているようです。
key は一引数をとる関数を指定し、リストのそれぞれの要素から比較キーを取り出すのに使います (例えば、 key=str.lower)。それぞれの項目に対応するキーは一度計算され、ソート処理全体に使われます。デフォルトの値 None は、別のキー値を計算せず、リストの値が直接ソートされることを意味します。（公式ドキュメント）
「tupleの先頭の要素でsortする」だけなら、key指定ありでもなしでも目的は果たせます。
ここで「じゃあkey指定なしでいいや〜」と思ってkey指定しないと、思わぬ罠にハマる可能性があります。
tuple同士の比較は遅いのです。
実験
「tuple同士の比較」 vs. 「tupleの要素同士の比較」
この2つで、どれほどの速度差があるのかを確認します。
実行環境はAtCoderのコードテスト（PyPy3）です。
各tuple内の要素数は3です。
import random
random.seed(1)
''' n: 3通り
N = 1*(10**3)
N = 3*(10**3)
N = 5*(10**3)
'''
tl = [] # list of tuples
for i in range(N):
l = random.randint(1,N)
r = random.randint(1,N)
tl.append((l,r,i)) # 3 elements in tuple
'''
# 1. tuple同士の比較
for i in range(N):
for j in range(N):
res = tl[i] < tl[j]
# 2. tupleの要素同士の比較
for i in range(N):
for j in range(N):
res = tl[i][0] < tl[j][0]
'''
計測結果は以下です。
| N | 1. tuple同士 | 2. tupleの要素同士 |
|---|---|---|
| $1*10^3$ | 78 ms | 17 ms |
| $3*10^3$ | 672 ms | 129 ms |
| $5*10^3$ | 1875 ms | 368 ms |
タプルの要素が3つなので「高々3倍程度の誤差か？」と思いきや、かなり大きな差が出ました。
怖い。
tuple配列のsort 「key指定なし」 vs. 「key指定あり」
import random
random.seed(1)
''' n: 3通り
N = 1*(10**5)
N = 3*(10**5)
N = 5*(10**5)
'''
tl = [] # list of tuples
for i in range(N):
l = random.randint(1,N)
r = random.randint(1,N)
tl.append((l,r,i)) # 3 elements in tuple
'''
3. tl.sort()
4. tl.sort(key=lambda x:x[0])
'''
測定結果は以下です。
| N | 3. keyなし | 4. keyあり |
|---|---|---|
| $1*10^5$ | 198 ms | 98 ms |
| $3*10^5$ | 735 ms | 359 ms |
| $5*10^5$ | 1361 ms | 708 ms |
最初の実験ほどではありませんが、およそ2倍程度の差が出ました。
おわりに
さほど気にする差ではないかもしれませんが、問題によってはこれでハマる可能性があります。
実験のコードを見て「あの問題かな？」と思った人もいるはず。
ちなみに、key指定はlambdaではなくitemgetterのほうが高速らしいですね。