【NumPy入門 np.where】条件にあったIndexの取得とif文的な使い方

こんにちは！インストラクターのフクロウです！

NumPyには、条件にあった配列の要素のインデックスを返してくれるnp.whereがあります。このnp.whereは要素のインデックスを返してくれるだけでなく、

などを設定することもできます。NumPyは高速な計算速度が売りなので、できればPythonのfor分やif文を使ってすべての要素を舐めるようなことはやりたくありません。

そんなとき、np.whereを使って代用できれば嬉しいですね。この記事でnp.where関数の使い方を抑えれば、すぐにプログラムに活かせますよ！

np.whereのAPI

np.where(
    condition,    # 条件式 
    x,            # conditionがTrueだったときに元の配列要素の代わりに代入される値 
    y             # conditionがFalseだったときに元の配列要素の代わりに代入される値
)

np.whereはnp.arrayを使った条件式を第一引数に指定します。第二引数、第三引数は省略でき、その場合は条件にあった要素のindexが返ってきます。第二引数、第三引数を指定する場合はどちらも指定します。

その場合、元の配列で条件式がTrueになった要素にxが、Falseになった要素にyが代入された配列が出力値になります。

np.whereの使い方

サンプル配列の準備

import numpy as np

a = np.arange(10)
b = a.reshape(2,-1)
c = np.arange(20).reshape(4,-1)
d = c.reshape(2,5,-1)

print("a")
print(a)

print("nb")
print(b)

print("nc")
print(c)

print("nd")
print(d)

[Output]

a
[0 1 2 3 4 5 6 7 8 9]

b
[[0 1 2 3 4]
 [5 6 7 8 9]]

c
[[ 0  1  2  3  4]
 [ 5  6  7  8  9]
 [10 11 12 13 14]
 [15 16 17 18 19]]

d
[[[ 0  1]
  [ 2  3]
  [ 4  5]
  [ 6  7]
  [ 8  9]]

 [[10 11]
  [12 13]
  [14 15]
  [16 17]
  [18 19]]]

さて、NumPyの配列に論理演算子を使うと、要素ごとにTrue/Falseが判定されるんでしたね。

実際に見てみましょう。（>>>はPythonのシェルで実行されたコードだと言うことを示しますが、上のコードも含めてjupyter labで試すことをおすすめします。）

>>> a%2==0

array([ True, False,  True, False,  True, False,  True, False,  True,
       False])

>>> b%2==0

array([[ True, False,  True, False,  True],
       [False,  True, False,  True, False]])

>>> c%2==0

array([[ True, False,  True, False,  True],
       [False,  True, False,  True, False],
       [ True, False,  True, False,  True],
       [False,  True, False,  True, False]])

>>> d%2==0

array([[[ True, False],
        [ True, False],
        [ True, False],
        [ True, False],
        [ True, False]],

       [[ True, False],
        [ True, False],
        [ True, False],
        [ True, False],
        [ True, False]]])

多次元配列であっても一次元配列と同様にTrue,Falseで埋まった配列ができていますね。np.whereの返り値とこれらを混同してしまうことが多々あるので、注意してくださいね。

条件を満たす要素のindexを取得

np.whereの最も基本的な使い方を見てみましょう。第一引数に条件式を指定するだけで、条件式がTrueになる要素のindexを得られます。

>>> np.where(a%2==0)

(array([0, 2, 4, 6, 8]),)

>>> np.where(b%2==0)

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

>>> np.where(c%2==0)

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

>>> np.where(d%2==0)

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

出力値に注目してください。出力はタプル型になっていて、元の配列の次元数分だけのnp.arrayを持っています。例えば、bで条件（要素が偶数である）に合う要素は

の5つだと読み取れます。

条件によって別要素を代入

次は少し発展的な用法です。第二引数にTrueだったときに代入したい値、第三引数にFalseだったときに代入したい値を指定してif分的に使いましょう。

>>> np.where(a%2==0, "偶数", "奇数") # Trueなら"偶数", Falseなら"奇数"を代入

array(['偶数', '奇数', '偶数', '奇数', '偶数', '奇数', '偶数', '奇数', '偶数', '奇数'],
      dtype='<U2')

>>> np.where(b%2==0, 1, 0) # Trueなら1, Falseなら0を代入

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

>>> np.where(c%2==0, c, c*100) # Trueならそのまま、 Falseなら「元の値*100」を代入

array([[   0,  100,    2,  300,    4],
       [ 500,    6,  700,    8,  900],
       [  10, 1100,   12, 1300,   14],
       [1500,   16, 1700,   18, 1900]])

>>> np.where(d%2==0, -d, d) # Trueならマイナスの符号をつけた値, Falseならそのままの値を代入

array([[[  0,   1],
        [ -2,   3],
        [ -4,   5],
        [ -6,   7],
        [ -8,   9]],

       [[-10,  11],
        [-12,  13],
        [-14,  15],
        [-16,  17],
        [-18,  19]]])

第二引数、第三引数には代入したい値を指定するんでした。もちろんその代わりに関数や式を書いてもOKです。これを応用すれば、様々な処理がnp.whereを使って実装できそうですね！

まとめ

この記事では、条件にあった要素のindexを返す関数np.whereを紹介しました。オプショナルなパラメータである第二引数と第三引数を指定することで、複雑な処理にも利用できそうな関数ですね。

この関数は案外むずかしい動作をするので、またわからなくなったらぜひまたこの記事を読んであげてください！