NumPyのソート(並び替え)のために知っておきたい5つの方法

2021 5/06

2020年10月6日 2021年5月6日

NumPyの配列のソートは、データ分析やデータエンジニアリングにおいて多用する操作の1つです。そこで、ここではそんな操作に役立つ5つの方法を解説します。その5つとは以下の通りです。

sort()
argsort()
lexsort()
partition()
argpartition()

まずsortとargsortはしっかりと使いこなせるようになっておいましょう。そして、その2つだけでは対応できない場合に、他の3つのソート方法を検討するようにしてみてください。きっとデータ分析の効率が上がるはずです。

それでは、それぞれ見ていきましょう。

はじめに: 配列の次元軸について

はじめに配列を扱う際に必須の知識である「次元軸」という概念について軽く復習しておきましょう。NumPyの配列は、1次元配列(ベクトル)、2次元配列(行列) 、3次元配列(テンソル)というように、「次元数」という属性を持っています。

以下のコードで、それぞれの配列を作成しています。

In [1]:

import numpy as np
# 1次元配列
a = np.arange(3)
print('1次元配列 \n', a, '\n')
# 2次元配列
b = np.arange(6).reshape(2, 3)
print('2次元配列 \n', b, '\n')
# 3次元配列
c = np.arange(12).reshape(2, 2, 3)
print('3次元配列 \n', c)

1次元配列 
 [0 1 2] 

2次元配列 
 [[0 1 2]
 [3 4 5]] 

3次元配列 
 [[[ 0  1  2]
  [ 3  4  5]]

 [[ 6  7  8]
  [ 9 10 11]]]

イメージで表すと、これらの配列のそれぞれの要素は以下の画像のように格納されています。

1次元配列は横軸のみです。2次元配列は縦軸(axis=0)と横軸(axis=1)の2つの軸で成り立っています。3次元配列は奥行き軸(axis=0)と縦軸(axis=1)、横軸(axis=2)の3つの軸で成り立っています。

NumPy配列のソートでは、この軸(axis)を指定して、どの方向に沿って要素を並べ替えるのかを選択することができるようになっています。それらをうまく使いこなすためにも、上の画像でそれぞれの次元数の配列で、axis=n がどの軸のことを指すのかをしっかり確認しておきましょう。最も高次元の軸から0, 1, 2 と数えていくというように覚えておけば便利です。

それでは、以下からNumPyの配列をソートするために使う各関数を見ていきたいと思います。

1. sort: 昇順ソートした配列を作成

np.sortは NumPy配列の要素を昇順ソートする関数です。第一引数にソートする配列、第二引数にソートする次元軸を渡します。

詳しくは『NumPyのsort関数で配列をソート(並び替え)する方法まとめ』で解説しているので、ここでは簡潔に見ていくことにします。

1次元配列のソート

以下のコードは1次元配列をソートしています。

In [1]:

import numpy as np
rng = np.random.default_rng()
a = rng.integers(0, 11, 5)
a

Out[1]:

array([ 9,  7,  8,  4, 10])

In [2]:

np.sort(a)

Out[2]:

array([ 4,  7,  8,  9, 10])

2次元配列のソート

多次元配列の場合はどうでしょうか。以下の配列を例に見てみましょう。

In [3]:

a = rng.integers(0, 30, (3, 5))
a

Out[3]:

array([[10, 28, 26, 14,  4],
       [ 2, 22, 13, 23, 10],
       [21,  9, 28,  1,  6]])

次元軸を指定しない場合は、1次元軸（横方向）にソートします。

In [4]:

# 1次元軸(横方向)にソート
np.sort(a)

Out[4]:

array([[ 4, 10, 14, 26, 28],
       [ 2, 10, 13, 22, 23],
       [ 1,  6,  9, 21, 28]])

別の次元軸を指定したい場合は、オプション引数 axis= を使います。

In [5]:

# 2次元軸（縦方向)にソート
np.sort(a, axis=0)

Out[5]:

array([[ 2,  9, 13,  1,  4],
       [10, 22, 26, 14,  6],
       [21, 28, 28, 23, 10]])

ndarray.sort()について
np.sort関数と同じ機能を持つものにndarray.sortメソッドというものもあります。こちらは、対象となる配列を引数に渡すのではなく、a.sort() というように先頭に書きます。そして、もう一つ重要な違いとして、np.sort関数は引数に渡した配列の要素をソートした新しい配列を作成しますが、ndarray.sortメソッドは、対象の配列そのものをソートするという違いがあります。目的に応じて使い分けると良いでしょう。

2. argsort: ソートに使うインダイスを取得

np.argsortは、配列をソートするために使用することができるインダイスを取得する関数です。

『NumPyの配列のスライスの必須テクニックまとめ』で解説している通り、NumPyの配列は、豊富なスライシングの方法が用意されています。そのため、このインダイスによって、np.sortではできないようなソート操作が可能になります。

詳しくは『NumPyのargsort関数で配列をソートしたインデックスを取得する方法』で解説していますので、ここでは簡潔に見ていくことにします。

np.sortと同じく第一引数に配列を、第二引数に次元軸を渡します。

1次元配列のソートインデックスの取得と使い方

まずは1次元配列を見てみましょう。以下の配列をご覧ください。

In [1]:

import numpy as np
rng = np.random.default_rng()
a = rng.integers(0, 11, 5)
a

Out[1]:

array([7, 1, 5, 8, 9])

この配列をargsortに渡すと以下のようにインダイスの配列を返します。

In [2]:

# ソートに使えるインダイスを取得
ind = np.argsort(a)
ind

Out[2]:

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

このインダイスで、元の配列をスライスすると、配列をソートすることができます。

In [3]:

# インダイスを使って並べ替え
a[ind]

Out[3]:

array([1, 5, 7, 8, 9])

2次元配列のソートインデックスの取得と使い方

多次元配列の場合、デフォルトでは最後の軸(axis=-1)である1次元軸（横方向）に沿ったソートインダイスを返します。以下の配列を使って確認しましょう。

In [1]:

import numpy as np
rng = np.random.default_rng()
a = rng.integers(0, 11, (2, 5))
a

Out[1]:

array([[2, 5, 1, 7, 0],
       [2, 2, 0, 6, 6]])

以下のように、横方向のインダイスを返します。

In [2]:

# ソートに使えるインダイスを取得
ind = np.argsort(a)
ind

Out[2]:

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

そして、多次元配列をインダイスでソートするにはtake_along_axis()を使います。

In [3]:

# インダイスを使って並べ替え
np.take_along_axis(a, ind, axis=-1)

Out[3]:

array([[0, 1, 2, 5, 7],
       [0, 2, 2, 6, 6]])

2次元軸（縦軸）でのソートインダイスを取得したい場合は、オプション引数axisで指定します。

In [4]:

# ソートに使えるインダイスを取得
ind = np.argsort(a, axis=0)
ind

Out[4]:

array([[0, 1, 1, 1, 0],
       [1, 0, 0, 0, 1]])

このインダイスを使って、配列を縦軸に沿って並べ替えてみましょう。

In [5]:

# インダイスを使って並べ替え
np.take_along_axis(a, ind, axis=0)

Out[5]:

array([[2, 2, 0, 6, 0],
       [2, 5, 1, 7, 6]])

3. lexsort: 複数のキーを設定してソートインダイスを取得

np.lexsortは、複数のキーでソートする関数です。取得するのはソートした配列ではなく、ソートインダイスです。「複数のキーでソート」というのは、基本的には配列b を基準に昇順ソートするが bの中に同値の要素があった場合は、次の配列a を基準にソートするというものです。

詳しくは『NumPyのlexsort関数で配列を複数のソートキーで並べ替える方法』で解説しているので、ここでは簡潔に解説します。

まずは以下のコードをご覧ください。

これは、まず名前で昇順にソートしていますが、名前が同じ場合は苗字を比較してソートするという処理になっています。

In [1]:

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

Out[1]:

array([2, 1, 0])

取得したインダイスで次のように確認してみると、意図した通りの並びになっていることがわかります。

In [2]:

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

Out[2]:

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

もう1つ見てみましょう。

以下のコードは、まず配列b で昇順ソートしており、b の要素が同値の場合は配列a で昇順ソートするというようになっています。

In [3]:

import numpy as np
a = np.array([5, 4, 3, 2, 1])
b = np.array([10, 20, 30, 40, 10 ])
lex_ind = np.lexsort((a, b))
lex_ind

Out[3]:

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

取得したインダイスを使って確認してみましょう。

In [4]:

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

Out[4]:

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

4. partition: 任意の仕切り値で要素を再配置

np.partitonは、配列の要素のうち、任意の値k を基準にして、それより小さい値を左側に、それより大きい値を右側に再配置する関数です。この関数もまた、sortやargsortなどではできないデータエンジニアリングを可能にしてくれます。

詳しくは『NumPyのpartition()で配列の要素を再配置する方法』で解説しているので、ここでは簡潔に見ていくことにしまう。第一引数に配列、第二引数に仕切り値となる値kthを指定します。

以下の配列を使って、np.partitionの使い方を確認していきます。

In [1]:

import numpy as np
a =np.array([ 5, 4, 1, 0, -1, -3, -4, 2])
a

Out[1]:

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

この配列a のkth(何番目に大きい値であるかの順番)は次のようになっています。

このように最小値から最大値の順に0から数えていきます。1から数えるのではない点にご注意ください。

以下のコードでは、np.paritionを使って、kth=1（=2番目に大きい値。この配列では-3が該当）で仕切って、それより小さい値を左側に、それより大きい値を右側に再配置しています。

In [2]:

np.partition(a, 1)

Out[2]:

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

これは以下の画像のように処理されています。

kth=1 より小さな値は左側に、大きな値は右側に再配置されています。

仕切り値kthは複数指定することが可能です。

In [3]:

np.partition(a, (2, 5))

Out[3]:

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

これは以下の画像のように処理されています。

ndarray.partition()について
np.partition関数と同じ機能を持つものにndarray.partitionメソッドというものもあります。こちらは、対象となる配列を引数に渡すのではなく、a.partition() というように先頭に書きます。そして、もう一つ重要な違いとして、np.partition関数は引数に渡した配列の要素を再配置した新しい配列を作成しますが、ndarray.partitionメソッドは、対象の配列そのものを再配置するという違いがあります。目的に応じて使い分けましょう。

5. argpartition: 任意の仕切り値での再配置インダイスを取得

np.argpartitionは、上述のpartitionと同じ容量でインダイスを取得する関数です。インダイスを活用することで、partitionよりも自由なソート操作が可能です。

繰り返しになりますが、インダイス配列を使ったデータ分析に熟達するためには『NumPyの配列のスライスの必須テクニックまとめ』で解説しているテクニックをまずはしっかりと吸収するようにしてください。

それでは、例によって『NumPyのargpartition()で配列の要素を再配置するインデックスを取得』で詳しく解説しているので、ここでは簡潔に見ていくこととします。

1次元配列のパーティション

以下の配列を例に見ていきましょう。

In [1]:

import numpy as np
a =np.array([ 5, 4, 1, 0, -1, -3, -4, 2])
a

Out[1]:

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

以下のコードでは、np.argparitionを使って、kth=1（=2番目に大きい値。この配列では-3が該当）で仕切って、それより小さい値を左側に、それより大きい値を右側に再配置した際のインダイスを取得しています。

In [2]:

ind = np.argpartition(a, 1)
ind

Out[2]:

array([6, 5, 2, 3, 4, 1, 0, 7])

このインダイスを使って、実際に配列の要素を再配置するには次のように書きます。

In [3]:

# インダイスを使って再配置
a[ind]

Out[3]:

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

2次元配列のパーティション

続いて2次元配列の場合を見てみましょう。以下の配列を例にします。

In [1]:

import numpy as np
rng = np.random.default_rng()
a = rng.integers(-30, 31, (3, 5))
a

Out[1]:

array([[ -6, -23, -11,  -4,   7],
       [-15,  -8, -29,   7, -11],
       [ -9,  18,  -2, -28,  23]])

argpartitionはデフォルトでは、最後の軸（横方向）でインダイスを取得します。

In [2]:

#  行ごとにインダイスを取得
ind = np.argpartition(a, 2)
ind

Out[2]:

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

多次元配列の場合で、このインダイスを使って、実際に配列を再配置するには、take_along_axis()を使います。

In [3]:

# 並べ替え
np.take_along_axis(a, ind, axis=-1)

Out[3]:

array([[-23, -11,  -6,  -4,   7],
       [-29, -15, -11,   7,  -8],
       [-28,  -9,  -2,  18,  23]])

オプション引数 axis で、再配置する軸を変更することが可能です。以下のコードでは2次元軸（縦軸）で再配置したインダイスを取得しています。

In [4]:

#  列ごとにインダイスを取得
ind = np.argpartition(a, 2, axis=0)
ind

Out[4]:

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

このインダイスを使って、take_along_axis() で配列を実際に再配置してみましょう。

In [5]:

#  並べ替え
np.take_along_axis(a, ind, axis=0)

Out[5]:

array([[-15, -23, -29, -28, -11],
       [ -9,  -8, -11,  -4,   7],
       [ -6,  18,  -2,   7,  23]])

6. まとめ

以上がNumPyの配列をソートする際に、頭に入れておきたい方法の全てです。最後にあらためて一覧を載せておきます。お役に立てば幸いです。

関数	説明
sort	配列をソートした新しい配列を作成。同名のメソッドあり。メソッドはオリジナルの配列を並び替える。
argsort	配列のソートインダイスを作成
lexsort	複数のキーによってソートインダイスを作成
partition	指定の仕切り値を境に、それより低い値を左側に、それより高い値を右側に再配置した新しい配列を作成。メソッドあり。メソッドはオリジナルの配列の要素を再配置する。
argpartition	指定の仕切り値を境に、それより低い値を左側に、それより高い値を右側に再配置するソートインダイスを作成。