私はpythonのnumpyを使ってよくテキストログを整形する機会がよくあるのですが、その際、１度つまずいたことのある内容について何度も調べることがあるのですが（１度解決したらどうでも良いやと思っている点は直したい…）、その際調べても出てきづらいnp.whereの返り値の使い方について備忘録として残しておきます。内容がチープなのは許してください…
##np.whereで特定の要素を含む行を取得
そもそものnp.whereの返り値はindexの情報を取得できます。そのことを２次元のnumpy配列を使って見ていきましょう。
test1.py
import numpy as np
array = np.array([["a","b","c"],["d","e","f"],[1,2,3]])
print(np.where(array == "d"))
実行結果
$ python test1.py
(array([1]), array([0]))
この実行結果は視覚的にはdという文字列は numpy 配列arrayのarray[1][0]にあるよ、ということを示しています。実際この返り値をどう扱うかは以下を見ていきましょう。
test2.py
import numpy as np
array = np.array([["a","b","c"],["d","e","f"],[1,2,3]])
print(np.where(array == "d"))
print(type(np.where(array == "d")))
print(np.where(array == "d")[0])
print(np.where(array == "d")[1])
print(np.where(array == "d")[0][0])
print(np.where(array == "d")[1][0])
実行結果
$ python test2.py
(array([1]), array([0]))
<class 'tuple'>
[1]
[0]
1
0
ということで np.where の返り値は tuple になっていて欲しい値によって適切にindexの指定をする必要があります。
今回は文字列dがある行(=1)を指定したいのでnp.where(array == "d")[0][0] を利用します。
test3.py
import numpy as np
array = np.array([["a","b","c"],["d","e","f"],[1,2,3]])
print(array[np.where(array == "d")[0][0]][:])
実行結果
$ python test3.py
['d' 'e' 'f']
##おまけ
np.where を使って色々試したことも書いておきます。
test4.py
import numpy as np
array = np.array([["a","b","c"],["d","e","f"],[1,2,3]])
print(array[np.where(array == "d")])
実行結果
$ python test4.py
['d']
numpy配列のindex指定の際に直接入れると指定した要素が配列に入った状態で返ってきます。
しかし、これでは、dという値を得られたわけではありません。
d自体を得るには以下のように、返ってきた配列のindexを指定しなければなりません。
test5.py
import numpy as np
array = np.array([["a","b","c"],["d","e","f"],[1,2,3]])
print(array[np.where(array == "d")][0])
実行結果
$ python test5.py
d
ここまでは、そもそも np.where を使う目的が特定の要素を含む index の取得であるのでこうした使い方はしねぇよって感じなんですが、おまけなのでご容赦ください…