NumPyのlexsort関数で配列を複数のソートキーで並べ替える方法

NumPyのlexsort関数は、複数のソートキー(ソートの基準)によって、配列をソートすることができる関数です。argsortと同じく戻り値はソートした要素のインダイスです。

引数に渡した配列、またはシーケンスのうち最後のキーから優先的に基準となりソートされ、同値があった場合は、最後から2つ目のキーを基準にしてソートします。1つの基準だけでは、意図した通りに配列の要素をソートできない場合に使います。

口で言うよりも実際に見た方がわかりやすいので、さっそく確認しておきましょう。なお、lexsortの前に『NumPyのargsort関数で配列をソートしたインデックスを取得する方法』を確認しておくことをおすすめします。

それでは始めましょう。

1. 書式

まずは書式を確認しましょう。

np.lexsort(keys, axis=-1)

2. サンプルコード

それではサンプルコードで確認していきましょう。

以下のコードでは、まず配列b をソートしています。そして配列b ではソートできない場合(値が同じ場合)に、配列a でソートしています。

import numpy as np
# lexsortでインダイスを取得(bでソートしてからaでソート)
a = np.array([5, 4, 3, 2, 1])
b = np.array([10, 20, 30, 40, 10 ])
lex_ind = np.lexsort((a, b))
lex_ind

array([4, 0, 1, 2, 3])

以下のコードをご確認ください。

まず配列b の要素が昇順で並び替えられています。そして、配列b の要素が同値の場合は、配列a を基準にして並び替えています。

[(a[i],b[i]) for i in lex_ind]

[(1, 10), (5, 10), (4, 20), (3, 30), (2, 40)]

わかりづらいと思いますので、np.argsortでbのみを基準としてソートした場合と比較してみましょう。

np.argsort(b)

array([0, 4, 1, 2, 3])

このインダイスにしたがって、配列a とb を並べ替えると、以下のようになります。

[(a[i],b[i]) for i in arg_ind]

[(5, 10), (1, 10), (4, 20), (3, 30), (2, 40)]

先ほどと見比べると、前2つの要素 (5, 10) と (1,10) の並びが異なっていることがわかります。これは lexsort では最初の基準であるb の要素が同値だった場合、次にa を基準に並び替えており、argsort では単純にb のみを安定オーダー(=同値の要素はインデックス番号順に並べ替えること）で並べ替えているからです。

このように、lexsort は、引数に渡したオブジェクトの内、後ろのものから順にソートしていき、それだけでは順序づけができない場合、後ろから2つ目の順にソートする、というようにして複数のソートキーを指定することができる関数です。

もう一つ例を見てみましょう。

以下の苗字と名前を格納した2つの配列をlexsortで並べ替えてみましょう。これは、まず名前で並べ替え、名前が同じ場合のみ苗字で並べ替えるという操作をしています。

import numpy as np
s = np.array(['やまだ', 'やまだ',  'あおき'])
f = np.array(['はじめ', 'いちろう', 'いちろう' ])
lex_ind = np.lexsort((s, f))
lex_ind

array([2, 1, 0])

取得したインダイスの順に並べ替えると次のようになります。

[s[i] + ' ' + f[i] for i in lex_ind]

['あおき いちろう', 'やまだ いちろう', 'やまだ はじめ']

次に、argsortで名前だけで並べ替えてみましょう。すると以下のようになります。

arg_ind = np.argsort(f)

[s[i] + ' ' + f[i] for i in arg_ind]

['やまだ いちろう', 'あおき いちろう', 'やまだ はじめ']

以上のように引数で渡した配列のうち、後ろにあるものを基準にソートし、同値の場合は後ろから2番目の配列を基準にソート、それでも同じなら後ろから3番目の配列を基準にソートというように、lexsortは複数の基準でのソートを可能にします。

3. まとめ

以上が、NumPyのlexsort関数です。よく見る関数ではないですが、配列の要素の並べ替えでハマってしまった時に、これを覚えておくと解決できる場合が多々あります。

以下の通常のソート方法とあわせて、頭に入れておくと良いでしょう。

• NumPyのsort関数で配列をソート(並び替え)する方法まとめ
• NumPyのargsort関数で配列をソートしたインデックスを取得する方法