【NumPy】配列の高速ソート

numpy.sort()

numpy.sort() は配列を受け取って、ソートされた配列の copy を返します（元の配列は変更されません）。

numpy.sort(array_like, axis=-1, kind='quicksort', order=None)

第 1 引数には配列に相当するオブジェクトを渡します。
疑似乱数で 1 次元配列を作ってソートしてみましょう。

# NUMPY_SORT

# In[1]

import numpy as np

# 乱数シードを71に設定
np.random.seed(71)

# 疑似乱数を使って配列を作成
# x = [8 9 9 3 2 1 7 2 4 1]
x = np.random.randint(1, 10, 10)

# 配列xをソート
y = np.sort(x)

print(y)
# [1 1 2 2 3 4 7 8 9 9]

多次元配列の場合、axis を指定しなければ、最後の軸に沿ってソートします。たとえば、2 次元配列のときは行ごとにソートします。これは axis=1 を指定した場合と同じ結果となります。

# In[2]

# 乱数のシードを11に設定
np.random.seed(11)

# 疑似乱数を使って配列を作成
# x = [[1 2 8 2]
#      [8 3 9 1]
#      [1 5 3 2]]
x = np.random.randint(1, 10, (3, 4))

# 配列をソート
y = np.sort(x)

print(y)
# [[1 2 2 8]
#  [1 3 8 9]
#  [1 2 3 5]]

2 次元配列で axis=0 を指定すると列ごとにソートします。
サンプルコード A で作った配列xを axis=0 でソートしてみます。

# In[3]

# 配列xを列ごとにソート
y = np.sort(x, axis=0)

print(y)
# [[1 2 3 1]
#  [1 3 8 2]
#  [8 5 9 2]]

kind 引数はソートする方法を quicksort, mergesort, heapsort から選択します。デフォルトの quicksort は O[NlogN] の速さをもつアルゴリズムです。特別な理由がない限り、この引数を指定する必要はありません。
 
構造化配列 をソートする場合は、order 引数はフィールド名を指定することができます。以下のサンプルコードでは、定期試験のデータ（氏名と英語・数学・国語の点数）を構造化配列に収めて、数学の成績順にソートします。

# NUMPY_SORT_FIELD

# In[1]

import numpy as np

# 複合データ型の定義
# U10 : 最大長10のUnicode文字列
# i1 : 1バイト整数(=np.int8)
d = [('名前', 'U10'), ('英語', 'i1'), ('数学', 'i1'), ('国語', 'i1')]

# 構造化配列を作成
score = np.zeros(3, dtype=d)

# 構造化配列へのデータの書き込み
score['名前'] = ['牛込友里', '滝野川小春', '比留間沙希']
score['英語'] = [72, 26, 81]
score['数学'] = [69, 15, 78]
score['国語'] = [92, 30, 83]

# 数学の成績順（降順）に配列をソート
x = np.sort(score, order='数学')

# 数学の成績順に名前を表示
print(x['名前'])
# ['滝野川小春' '牛込友里' '比留間沙希']

ndarray.sort()

コピーを作らずに配列自体をソートして変更する場合は ndarray.sort()メソッドを使います。

ndarray.sort(axis=-1, kind='quicksort', order=None)

使い方は np.sort() 関数と同じです。

# NUMPY_NDARRAY_SORT

# In[1]

import numpy as np

# 乱数シードを19に設定
np.random.seed(19)

# 疑似乱数を使って配列を作成
#[[6 3 9 4 5]
# [3 9 7 3 6]
# [2 7 3 8 8]
# [1 8 8 5 6]]
x = np.random.randint(1, 10, (4, 5))

# 列ごとにソート
y = x.sort(axis=0)

print(x)
# [[1 3 3 3 5]
#  [2 7 7 4 6]
#  [3 8 8 5 6]
#  [6 9 9 8 8]]

numpy.argsort()

numpy.argsort() はソートされた要素のインデックスを返します。

numpy.argsort(a, axis=-1, kind='quicksort', order=None)

引数の指定方法は numpy.sort() と同じです。

# NUMPY_ARGSORT

# In[1]

import numpy as np

# 乱数シードを20に設定
np.random.seed(20)

# 疑似乱数を使って配列を作成
#[[4 5 7 8 3]
# [1 7 9 6 4]
# [1 7 7 1 6]
# [8 6 3 7 4]]
x = np.random.randint(1, 10, (4, 5))

# 行ごとにソートして要素のインデックスを表示
y = np.argsort(x)

print(y)
# [[4 0 1 2 3]
#  [0 4 3 1 2]
#  [0 3 4 1 2]
#  [2 4 1 3 0]]