NumPyで条件に合う要素のインダイスを取得するargwhere()の使い方

NumPyのargwhere()関数は、① 0以外の要素のインダイスを要素ごとに取得します。または、② 指定の条件に合致する要素のインダイスを要素ごとに取得します。

①と同じ働きをする関数には、ほかにnumpy.nonzero()もあります。しかし、こちらは0以外の要素のインダイスを次元軸ごとに取得します。結論から言えば、配列の高度なスライスを行うにはnumpy.nonzero()の方が適しています。

②についても、指定の条件に合致する要素を操作するには、numpy.where()関数の方が適しています。

そのため、numpy.argwhere()の使用頻度は高くありません。このページでは、上述の関数も含めて簡潔に解説していきます。

1. numpy.argwhereの使い方

まずは書式を確認しましょう。とは言っても、この関数にはオプション引数はありません。

np.argwhere(a)

Note

np.argwhere(a)は、np.transpose(np.nonzero(a))とほとんど同じです。しかし、こちらの方が0次元の配列(要素が1つの配列：スカラー)に対しても正しく機能します。ただし、配列のスライシングにおいては、argwhere()で取得するインダイスよりもnonzero()で取得するインダイスの方が適しています。

2. サンプルコード

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

まずは以下の2次元配列a を作成します。

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

array([[3, 0, 0],
       [0, 4, 0],
       [5, 6, 0]])

この配列をnumpy.argwhere()に渡すと、値が0以外の要素のインダイスをインダイスごとに返します。

# 値が0以外の要素のインダイスを要素ごとに取得
np.argwhere(a)

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

このインダイスを使って値を取り出してみましょう。

# 要素の取得
for i in np.argwhere(a):
    print(a[tuple(i)])

3
4
5
6

なお、配列のスライスにはnumpy.nonzero()の方が遥かに適しています。この関数は、以下のコードのように、0以外の要素のインダイスを次元軸区切りで取得します。

#  配列のスライシングにはnp.nonzero()の方が適しています
ind = np.nonzero(a)
ind

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

この形のインダイスであれば、直接、配列をスライスすることで値を取得することができます。

a[ind]

array([3, 4, 5, 6])

このため、配列のスライス目的でのインダイスの取得ならば、numpy.nonzero()を使います。

一方で、numpy.argwhere()の利点は条件式定義の自由度にあります。numpy.nonzero()は0以外の要素のインダイスしか取得できませんが、numpy.argwhere()は条件式を書くことで、様々な値のインダイスを取得することが可能です。

#  np.argwhere()の利点は条件式定義の自由度にあります
np.argwhere(a >= 5 )

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

なお、条件に合う要素、合わない要素を任意の値に置換したい場合は、numpy.where()を使います。

# 条件に合う要素を100に、合わない要素を0に置換
np.where(a>=5, 100, 0)

array([[  0,   0,   0],
       [  0,   0,   0],
       [100, 100,   0]])

3. まとめ

以上がnumpy.argwhere()の使い方です。