NumPyのargpartition()で配列の要素を再配置するインデックスを取得

2021 5/06

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

NumPyのargpartiton()は、配列の要素のうち、任意の値k を基準にして、それより小さい値を左側に、それより大きい値を右側に再配置した際のインダイスを取得する関数です。”partition” は「仕切り」という意味です。つまり、配列を任意の値k で仕切って再配置するという意味です。

この関数で取得するインダイスの配列を使うと、配列のデータを様々に操作することができるようになります。このページでは、その点も含めて解説していきます。

なお、『NumPyのpartition()で配列の要素を再配置する方法』もあわせてご覧いただくと理解がより深まると思いますので、ぜひ参考にして頂ければと思います。

それではさっそく見ていきましょう。

NumPy配列のソート方法のまとめ
NumPyの配列をソートする上で知っておきたい5つの関数については『NumPyのソート(並び替え)のために知っておきたい5つの方法』でまとめて簡潔に解説しています。ぜひそちらもご確認ください。

1. 書式

まずは書式を確認します。

np.partition関数

書き方:

numpy.argpartition(a, kth, axis=-1, kind='introselect', order=None)

パラメーター:

引数	型	解説
a	array_like	ソートする配列を渡します。
kth	int または sequence of ints	配列の再配置の基準とする要素(仕切りとなる要素)を指定します。最小値を0として数えます。そのため、例えば配列の要素の中で2番目に小さな値を基準にする場合はkth=1になります。指定した要素は普通にソートした時と同じ位置に配置され、それより値が小さい要素は前に、値が大きい要素は後ろに配置されます。複数指定することも可能です。
axis*	int or None	要素を再配置する次元軸を渡します。Noneの場合はソートの前に1次元配列化されます。デフォルトは-1です。この場合は最後の軸に沿ってソートします。
kind*	{‘introselect’}	ソートのアルゴリズムです。使うことはほとんどありません。
order*	str or list of str	構造化配列を再配置する場合に使います。基準とするフィールドを指定します。
* はオプション引数であることを示します。

戻り値:

index_array：ndarray
配列aの要素をkthで仕切った際の配列のインダイスを返します。a が1次元配列の場合、a[index_array]と書くことで、ある配列の要素を実際に再配置することができます。より一般的には、np.take_along_axis(a, index_array, axis=a)と書くと元の配列の次元数に関係なく配列の要素を再配置することが可能です。

一緒に確認したい関数:

partition: 配列の要素を再配置
sort: 配列の要素をソート
argsort: 配列の要素をソートしたインダイスを取得

2. サンプルコード

それではサンプルコードを見ていきましょう。なおオプション引数のkind やorder は使うことはほとんどないので、ここでは割愛します。

1次元配列の再配置インダイスを取得

それでは以下の配列を使って、np.argpartitionの使い方を確認していきましょう。なお、先に『NumPyの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.argparitionを使って、kth=1（=2番目に大きい値。この配列では-3が該当）で仕切って、それより小さい値を左側に、それより大きい値を右側に再配置した際のインダイスを取得しています。

In [2]:

ind = np.argpartition(a, 1)
ind

Out[2]:

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

このインダイスを使って、実際に配列の要素を再配置するには、以下のコードのように配列をインダイスでスライスします。

なお、スライスのテクニックは、データサイエンスにおいて必須です。『NumPyの配列のスライスの必須テクニックまとめ』では、重要なテクニックをすべてまとめていますので、ぜひご確認ください。

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]])

問題

ここで少し複雑な操作を考えてみましょう。

np.argpartitionで取得したインダイスの配列を使うと、特定の行や列の要素のみを再配置したり、特定の行や列のインダイスを配列内の全行または全列に適用するというようなnp.partition関数ではできない操作を行うことが可能です。

それでは例えば、特定の行のインダイスを、配列の全ての行に適用するにはどうすればいいでしょうか？また、特定の列のインダイスを、配列の全ての列に適用するにはどうすればいいでしょうか？

次の配列を例に考えてみましょう。

In [1]:

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

Out[1]:

array([[ -3,   9,   8, -24,  22],
       [ 26, -12, -18,  -4, -11],
       [ 14, -28,   3,  11,  23]])

解答は次の通りです。

特定の行のインダイスを全ての行に適用するには次のように書きます。

In [2]:

# 真ん中の行の再配置インダイスを取得
ind = np.argpartition(a[1], 2)

In [3]:

#  真ん中の行のインダイスを全行に適用
a[:, ind]

Out[3]:

array([[  8,   9,  22, -24,  -3],
       [-18, -12, -11,  -4,  26],
       [  3, -28,  23,  11,  14]])

特定の列のインダイスを全ての列に適用するには次のように書きます。

In [4]:

#  最初の列の再配置インダイスを取得
ind = np.argpartition(a, 1, axis=0)[:, 0]

In [5]:

# 最初の列のインダイスを全列に適用
a[ind]

Out[5]:

array([[ -3,   9,   8, -24,  22],
       [ 14, -28,   3,  11,  23],
       [ 26, -12, -18,  -4, -11]])

繰り返しになりますが、『NumPyの配列のスライス（要素の抽出）に必須の基本テクニックまとめ』で解説している通り、NumPyの配列には、リストやタプルよりも遥かに豊富なスライシングのテクニックがあります。データエンジニアリングで、データを自由に操作するには必須のスキルなので、何度でもご確認ください。