Numpy.where で配列から条件で要素を抽出する

ゆうすけ

Numpy の where 関数の使い方を知りたいです。

資格マフィア

サンプルコードを用いて解説しよう。
配列に条件を適用することができるぞ。

✔️ 本記事のテーマ

　Numpy.where 関数の使い方

✔️ 読者さんへの前置きメッセージ

本記事は
「Numpy.where 関数の使い方」　や　「配列から条件に適した要素を抽出する方法」
について書いています。

Numpy.where 関数を使うことで、配列に条件を適用することができます。

この記事に掲載しているサンプルコードは環境さえ整っていれば、
コピペで動くはずなので使用して頂いてOKです。

では、解説していきましょう。

Numpy.where 関数の使い方
Numpy.where 関数で複数の条件を適用する
Numpy.where 関数で条件を満たす場合だけ置換する
Numpy.where 関数で条件を満たすインデックスを取得する
Numpyのスキルをつけるために最適な一冊

Numpy.where 関数の使い方

numpy.where() は引数に条件, x, y を指定して使います。

対象の ndarray に対して、
条件を満たす（True）場合に x、満たさない（False）場合に y へと
置換した ndarray を返します。

import numpy as np


before = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])

after = np.where(before < 5, 0, 1)

print(after)
# [0, 0, 0, 0, 1, 1, 1 ,1 ,1]

上記のコードでは、条件(=5未満)を満たす場合に 0 , 満たさない場合に 1 に置換しています。

なお、条件は before < 5 というように ndarray を対象として記述します。

このように記述することで before の各要素に対して、
条件を満たす/満たさないの判定と置換を実施します。

Numpy.where 関数で複数の条件を適用する

numpy.where() では複数の条件を適用することができます。

ただ、複数条件を記述する場合は Numpy 独自のルールに従う必要があります。

具体的には、各条件式を () で囲み、& や | を使います。

before = np.array([-3, -2, -1, 1, 2, 3, 7, 8, 9])

after = np.where((0 < before) & (before < 5), 0, 1)

print(after)
# [1, 1, 1, 0, 0, 0, 1 ,1 ,1]

なお、条件式に and , or を使用すると実行時に error が発生します。

ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()

Numpy.where 関数で条件を満たす場合だけ置換する

numpy.where() では条件を満たす要素だけを置換することもできます。

引数 y に処理対象の ndarray を渡すことで、
条件を満たす要素は引数 x で指定した値に、満たさない要素はそのままにすることができます。

import numpy as np


before = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])

after = np.where(before < 5, 0, before)

print(after)
# [0, 0, 0, 0, 5, 6, 7, 8, 9]

ちなみに、引数 x に処理対象の ndarray を渡すことで、
逆に、条件を満たさない要素だけを置換することもできます。

import numpy as np


before = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])

after = np.where(before < 5, before, 1)

print(after)
# [1, 2, 3, 4, 1, 1, 1 ,1 ,1]

Numpy.where 関数で条件を満たすインデックスを取得する

numpy.where() の引数として条件だけを渡した場合（置換先を指定しなかった場合）は、
条件を満たす要素のインデックス（位置）を返します。

各次元（行、列）に対して条件を満たすインデックス（行番号、列番号）の ndarray のタプルを返します。

注意点として、1次元配列に対して Numpy.where 関数をかけても、タプルを返します。

before = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])

after = np.where(before < 5)

print(after)
# (array([0, 1, 2, 3]),)

このインデックスと for を使うことで、
条件を満たす要素にだけ特定の処理をかけることができます。

before = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9])

for ix in np.where(before < 5)[0]:
    before[ix] *= 100

print(before)
# [100 200 300 400   5   6   7   8   9]

np.where(before < 5)[0] とすることで条件を満たすインデックスのイテレータを取り出しています。