【python】numpy配列の特定の条件を満たす要素のインデックスを取得する

メモ。

numpy.arrayオブジェクトの要素のうち、指定した条件を満たす要素は下記のように取得できる。

import numpy as np

arr = np.array([[1,2,3],[4,5,6]]) # array([[1, 2, 3],
                                  #       [4, 5, 6]])

# 特定の条件を満たす要素を取得
arr[arr > 3] # array([4, 5, 6])

# 特定の条件を満たす要素を取得
arr[arr > 3] # array([4, 5, 6])

ただし、指定した条件を満たす要素の「添字」を取得したいときは、numpy.argwhereを使うと良い。

np.argwhere(arr > 3) # array([[1, 0],
                     #            [1, 1],
                     #            [1, 2]])

numpy配列のshapeに合わせた添字を返してくれるので、便利。