NumPy where どの列を調べるか指定する

前回に引き続き、NumPyのちょっと便利な使い方です。

Numpy.where

は配列の中の要素を調べるための有名な関数で、皆さん使ってるかと思います。

標準的な使い方は下記の通りです。

import numpy as np

pieces = np.array([
        [1, 0, 0],
        [0, 0, 1],
        [0, 0, 0]
        ])

array = np.array([1,0,0])

ndarray = np.where(array == 1, True, False)

print(ndarray)

#下記が出力されます。
[ True False False]

下記のサイトさんなどにも詳しいです。

NumPyで条件に応じた処理を行うwhereの使い方

他のサイトさんであまり紹介されていなくって、今回書いておきたいwhereの使い方が、2次元配列をwhereで調べる時に、列を指定する、というやり方です。

import numpy as np

pieces = np.array([
        [1, 0, 0],
        [0, 0, 1],
        [0, 0, 0]
        ])

ndarray = np.where(pieces == 1, True, False)[1]

print(ndarray)

#下記が出力されます。
[False False  True]

上記のように、whereの後にインデックス[1]をつけると、piecesの2列目（0からカウントなので、1は2番目）が1かどうかを判定して、True, Falseを返してくれてます。