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

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

LINEで送る
Pocket

ITエンジニアへ転職したい方におすすめ

自分を評価してくれる企業に転職して年収を上げたい! 自分のスキルにあった独自案件を知りたい!
エンジニアは今もっとも注目されている職業の1つ。エンジニアになって年収を増やしたい方や、あなたのスキルに見合った企業へ転職したいエンジニアの方も多いですよね。

しかし、大手の転職媒体は扱う求人数が多くても、誰もが登録しているので競争率もかなり高くなっています。そのため、あなたの条件に見合った企業を見つけても転職するためには、相応の努力とスキルが必要となります。

こういった媒体では、未経験からエンジニアを目指す方やエンジニア歴2〜3年で転職を考えている方にとって、最適な転職環境とはいえません。

そこでオススメしたいのが、未経験者や若手エンジニア向けの独自案件を多く掲載している「侍ワークス」です。

侍ワークスは、独自案件を多く掲載しているだけでなく、

・応募から就業まで一貫したサポート

・就業後もアフターフォロー

といった経験の浅い方や初めてエンジニアを目指す方にも安心のフォロー体制が整っています。もちろん登録は完全無料!しかも案件を見るだけなら登録も不要です。

まずは、お気軽にどんな求人があるか見てみてください。あなたにピッタリの企業がきっと見つかりますよ! 侍ワークスの求人情報を見る

書いた人

フクロウ

フクロウ

第一言語はPythonです。
皆さんRustやりましょう。