NumPyのargsort関数で配列をソートしたインデックスを取得する方法

np.argsort関数は、配列の要素をソートした時のインダイス（インデックスの複数形）を返す関数です。細かい部分はnp.sort関数と同じですが、配列をソートするのではなくインダイスを作成する点が異なります。

一見、np.sort関数があれば、np.argsort関数は不要に思えるかもしれません。しかしNumPyの配列には、Python標準のリストやタプルよりも豊富なスライシングのテクニックが実装されています。そのため、インダイスの配列を取得することで、np.sort関数では不可能な操作も行えるようになります。

そのために、ここではnp.argsort関数について解説していきます。以下のページもあわせて読んで頂くことで、理解が深まるでしょう。

• 『NumPyの配列のスライスの必須テクニックまとめ』
• 『numpy.sort関数で配列をソート(並べ替え)する方法まとめ』

それでは始めましょう。

1. 書式

早速、書き方を確認しましょう。

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

2. サンプルコード

それでは、サンプルコードを見ながら使い方を確認していきましょう。

1次元配列のソート

np.argsort関数は、配列の要素をソートするのではなく、ソートした時のインデックスを返します。以下の配列を例に確認しましょう。

import numpy as np
arr = np.array([2, 3, 1])
arr

array([2, 3, 1])

この配列を、np.argsort関数に渡すと次のように返ってきます。

ind = np.argsort(arr)
ind

array([2, 0, 1])

これは、配列をソートしたとして、元の配列の要素がどのように並べられるかを示すインダイス（インデックスの複数形）です。

ちょうど、下図のようになっています。

取得したインダイスを使ってスライスすることで、配列をソートすることができます。

#  インダイスでスライスして並べ替え
arr[ind]

array([1, 2, 3])

多次元配列のソート

多次元配列の場合、デフォルトのaxis=-1では横方向にソートした時のインダイスを返します。

import numpy as np
rng = np.random.default_rng()
x = rng.integers(0, 10, (2, 3))
x

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

# デフォルトの軸はaxis=-1(横方向)
ind = np.argsort(x)
ind

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

多次元配列をインダイスでソートする場合は、np.take_along_axis関数を使います。

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

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

次元軸を指定してソート

デフォルトの横方向以外の軸で、ソートしたい場合は、オプション引数 axis= で任意の軸を指定します。2次元配列の場合は、axis=0 で縦軸でソートすることができます。3次元配列の場合は、axis=0 で奥行き軸、axis=1 で縦軸でソートすることができます。この点は下図をご覧いただくとわかりやすいでしょう。

配列の次元軸について混乱してしまう場合は、『numpy.sort関数で配列をソート(並べ替え)する方法まとめ』ではより詳しく解説しているので、そちらをご覧頂ければと思います。

それでは、上で作成した2次元配列 x を、縦軸でソートしてみましょう。

# 2次元配列を縦軸(axis=0)でソート
ind = np.argsort(x, axis=0)
ind

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

この場合、np.take_along_axis関数で元の配列をソートする際も axis=0 を指定します。

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

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

ソートアルゴリズムを選択

オプション引数の kind= で、ソートアルゴリズムを選択することができます。選択できるソートアルゴリズムは、以下があります。

• quicksort（クイックソート）
• mergesort（マージソート）
• heapsort（ヒープソート）
• stable（安定ソートのどちらか）

それぞれのアルゴリズムについて詳しくは『numpy.sort関数で配列をソート(並べ替え)する方法まとめ』で解説していますので、ぜひご確認ください。通常は、デフォルトのクイックソートで問題ありません。

# ソートアルゴリズムの指定
np.argsort(x, kind='mergesort')

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

構造化配列のソート

構造化配列(Structured Array)をソートする場合は、オプション引数 order= で、どのデータを対象にソートするかを選択することができます。

構造化配列とは、以下のように、複数のデータ型をもつ配列のことです。詳しくは、『NumPyのdtype属性の一覧と参照・指定・変更方法』で解説しています。

import numpy as np
dtype = [('name', 'S6'), ('height', float), ('age', int)]
values = [('Taro', 1.8, 41), ('Jiro', 1.9, 38), ('Saburo', 1.7, 38)]
x = np.array(values, dtype=dtype)
x

array([(b'Taro', 1.8, 41), (b'Jiro', 1.9, 38), (b'Saburo', 1.7, 38)],
      dtype=[('name', 'S6'), ('height', '<f8'), ('age', '<i8')])

この配列には、文字列型の’name’・float型の’height’・int型の’age’という異なるデータ型のデータがあります。オプション引数 order= で、どのデータを基準にソートするかを指定します。

# heightを基準としてソート
np.argsort(x, order='height')

array([2, 0, 1])

# age
np.argsort(x, order='age')

array([1, 2, 0])

もし、ageが同じ場合は、heightでソートするというようにしたい場合は、次のように、order=['age', 'height'] と書きます。

np.argsort(x, order=['age', 'height'])

array([2, 0, 1])

3. インダイスを使ったソートのテクニック

NumPyの配列は、Python標準のリストやタプルよりも遥かに豊富なスライシングのテクニックがあります。

そのため、np.argsort関数で取得したインダイスの配列を使うと、特定の行や列の要素のみを並べ替えるというようなnp.sort関数ではできない操作を行うことが可能です。

それらの操作を行えるようになるには、『NumPyの配列のスライス（要素の抽出）に必須の基本テクニックまとめ』で解説しているように、基本的なスライスとインダイスの配列を使ったスライスをマスターしておく必要があります。

ここでは、この中のテクニックの1つを使って、2次元配列を特定の行、または列を基準にソートしてみましょう。

それでは、次の配列を例に見ていきます。

import numpy as np
rng = np.random.default_rng()
x =rng.integers(0, 10, (3, 5))
x

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

特定の行を基準に並び替え

この配列 x を使って、1行目を基準に並び替えてみましょう。まずは1行目をソートしたインダイスを取得します。

# 任意の行（1行目)のインダイスを取得
ind = np.argsort(x[0])
ind

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

取得したインダイスで次のようにスライスすると、1行目を基準に並べ替えられます。

# インダイスを使って特定の行を基準に並び替え
x[:, ind]

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

特定の列を基準に並び替え

特定の列を基準とした並び替えも方法は同じです。まずは、任意の列（ここでは3列目）のインダイスを取得します。

# 任意の列(3列目)のインダイスを取得
ind = np.argsort(x[:, 2])
ind

array([1, 2, 0])

このインダイスでスライスすると、3列目を基準に要素を並び替えることができます。

# インダイスを使って特定の列を基準に並び替え
x[ind]

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

4. まとめ

以上が、np.argsort関数の使い方です。繰り返しになりますが、この関数を使いこなすためにも、『NumPyの配列のスライスの必須テクニックまとめ』をしっかりと読み込んでおきましょう。