【初心者向け】ある条件式を満たす要素のインデックスを取得したい

背景

プロpythonistからすれば普通のことなのかもしれないんですが、
自分は理解するのに時間がかかり、
理解できたらその強さに感動した。。。ので。。。

ndarray[条件式]

行列のままで各要素がある条件を満たすか満たさないかを判別し、
その結果のTrue/Falseを同じ形状の行列として出力してくれます。

例

arr = np.array([[0,1,2,3],[0,2,4,6]])
print(arr<3)
# [[ True  True  True False]
#  [ True  True False False]]

numpy.nonzero

行列の0でない要素のインデックスを取得してくれ、x,y別のndarray配列として出力されます。

# 例の定義
arr_int = np.array([[3,5,0],[0,4,0]])
arr_bool = np.array([[True,True,False],[False,True, False]])

# np.nonzeroの利用
nonzero_int_row, nonzero_int_column = np.nonzero(arr_int)
nonzero_bool_row, nonzero_bool_column = np.nonzero(arr_bool)

# 各変数の値
# row: array([0, 0, 1])
# column: array([0, 1, 1]))

この例では、0でないのがarr[0,0],arr[0,1],arr[1,1]なのでこういう結果になります。

合体！

arr = np.array([[0,1,2,3],[0,2,4,6]])
arr_bool = arr<3
nonzero_row, nonzero_column = np.nonzero(arr_bool)
# row: [0 0 0 1 1]
# column: [0 1 2 0 1]

コメント

コメントいただきました！
np.whereが同じ機能で便利ということです！
僕が見てたサイトはなぜこれを使わなかったんだろう…