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

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

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

  • 条件文がTrueになる要素にする操作
  • 条件分がFalseになる要素にする操作

などを設定することもできます。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で条件(要素が偶数である)に合う要素は

  • b[0,0]
  • b[0,2]
  • b[0,4]
  • b[1,1]
  • b[1,3]

の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を紹介しました。オプショナルなパラメータである第二引数と第三引数を指定することで、複雑な処理にも利用できそうな関数ですね。

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

この記事を書いた人

【プロフィール】
DX認定取得事業者に選定されている株式会社SAMURAIのマーケティング・コミュニケーション部が運営。「質の高いIT教育を、すべての人に」をミッションに、IT・プログラミングを学び始めた初学者の方に向け記事を執筆。
累計指導者数4万5,000名以上のプログラミングスクール「侍エンジニア」、累計登録者数1万8,000人以上のオンライン学習サービス「侍テラコヤ」で扱う教材開発のノウハウ、2013年の創業から運営で得た知見に基づき、記事の執筆だけでなく編集・監修も担当しています。
【専門分野】
IT/Web開発/AI・ロボット開発/インフラ開発/ゲーム開発/AI/Webデザイン

目次