NumPy配列の条件に合った要素を置換するwhere関数の使い方

numpy.whereは、NumPy配列の要素のうち、特定の条件に合致するものをx で置換し、それ以外をy で置換する関数です。

より厳密には numpy.where()は2つの使い方が可能です。

1. 条件に合致する要素のインデックスを取得する。
2. 条件に合致する要素をxに、条件に合致しない要素をyに置換する。

ただし、1. については、numpy.where()よりも、numpy.nonzero()を使うことが推奨されています。こちらの方が、NumPy配列のサブクラス（マトリックスなど）も適切に処理してくれるからです。そのため、numpy.where()の用法は基本的に 2. になります。

ここでは、このnumpy.whereについて解説します。

1. numpy.whereの使い方

それでは実際のコードで確認していきましょう。その前に、まずは基本的な書き方を見ておきましょう。

numpy.where(condition[, x, y])

Note

引数に渡す全ての配列が1次元の場合、np.where()は以下と同じです。

[xv if c else yv for c, xv, yv in zip(condition, x, y)]

2. サンプルコード

それでは早速、実際のコードで使い方を確認していきましょう。

なお冒頭で述べたとおり、numpy.where()の用法は「配列の中の条件に合致する要素をxに、条件に合致しない要素をyに置換すること」です。しかし、それを使いこなすためには条件式を自在に操作できる必要があります。そのため、このページでは条件式についての解説もしっかりと行うために、あえてインデックスを取得する用法についても細かく見ていくことにします。

それでは早速はじめましょう。

2.1. 条件に合う要素のインデックスを取得

numpy.where()に、条件式(condition)だけを渡すと、その条件に合致する要素のインデックスをタプル型で返します。

1次元配列、2次元配列, … N次元配列でインデックスの返し方が異なるので、それぞれ確認していきましょう。

1次元配列の場合

まずは以下の1次元配列を使って見ていきましょう。

import numpy as np
#  配列を生成
arr1 = np.array([1, 9, 2, 3, 6, 8])
print(arr1)

[1 9 2 3 6 8]

この配列の中から、5より小さな要素を探したい場合は、次のように書きます。

#  5より小さな要素のインデックスを取得
np.where(arr1 < 5)

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

戻り値はタプル型です。そして、array()内の数値は、条件である「5より小さな要素」のインデックス番号を表わします。つまり、インデックス0と2と3が、指定した条件に合致する要素が存在する場所だということです。

実際に、関数から返ってきたインデックス番号でスライスすると数値を確認することができます。

#  関数が返したインデックスでスライスしてみましょう。
print(arr1[0])
print(arr1[2])
print(arr1[3])

1
2
3

このようにnp.where()は、配列の中の条件に合致する要素の位置を返す関数です。

なお、わざわざ上のように書かなくても、np.where関数の戻り値でスライスすると条件に合致した要素を取り出した新たな配列を生成することができます。

arr1[np.where(arr1 < 5)]

array([1, 2, 3])

他の条件式でも見てみましょう。

以下のコードでは、偶数のインデックスと、奇数のインデックスを取得しています。

#  偶数のインデックスを取得
even = np.where(arr1%2 == 0)
print(even)

#  奇数のインデックスを取得
odd = np.where(arr1%2 == 1)
print(odd)

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

2次元配列の場合

2次元配列の場合は、条件に合致する要素のインデックスをマトリックスで返します。

以下のコードをご覧ください。

#  2次元配列の場合
arr2 = arr1.reshape(2, 3)
print(arr2, '\n')

#  5より小さな要素のインデックスを取得
print(np.where(arr2 <5))

[[1 9 2]
 [3 6 8]] 

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

今度は配列が2つ入ったタプルが返ってきました。最初の配列は2次元軸（縦軸）方向のインデックスを、次の配列は1次元軸（横軸）方向のインデックスを表しています（Pythonのインデックス番号は0から始まります）。

そのため、この戻り値は、arr2[0, 0]、arr2[0, 2]、arr2[1, 0]に、条件(arr2 < 5)に合致する要素があるということを示しています。

print(arr2[0, 0])
print(arr2[0, 2])
print(arr2[1, 0])

1
2
3

2.2. 複数の条件を指定

一つひとつの条件式を丸括弧()で囲み、以下のビット演算子を使うと、複数条件を指定することができます。

なお、ビット演算子を含む演算子は『Pythonの演算子の一覧表と解説』で解説しています。

それでは、以下の配列を例に、実際のコードで確認していきましょう。

import numpy as np
arr4 = np.array([5, 2, 1, 3, 9, 11, 10, 7, 12, 4, 6, 8])

複数条件を同時に満たす要素のインデックスを取得（&演算子）

以下のコードでは、&演算子を使って、「2以上」「5以下」の条件を両方満たす要素のインデックスを取得しています。

#  論理積：2以上5以下の要素のインデックスを取得
np.where((arr4 >= 2 ) & (arr4 <= 5))

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

実際にスライスで確認してみると、確かに条件を満たしている要素のインデックスを返していることがわかります。

arr4[np.where((arr4 >= 2 ) & (arr4 <= 5))]

array([5, 2, 3, 4])

複数条件を1つ以上満たす要素のインデックスを取得(|演算子)

|演算子を使うと、複数の条件をどれか1つでも満たす要素のインデックスを取得することができます（複数の条件を同時に満たすものも含む）。

以下のコードでは、「偶数である」「9である」のどちらか一方でも満たす要素のインデックスを取得しています。

#  論理和: 偶数と9のインデックスを取得
np.where((arr4%2==0) | (arr4 ==9))

(array([ 1,  4,  6,  8,  9, 10, 11]),)

スライスで実際に確認してみましょう。

確かに条件を満たす要素のインデックスであることがわかります。

arr4[np.where((arr4%2==0) | (arr4 ==9))]

array([ 2,  9, 10, 12,  4,  6,  8])

なお、複数の演算子を同時に使うこともできます。

以下のコードでは、「2以上5以下」という条件と、「11である」という条件のどちらか一方でも満たす要素のインデックスを取得しています。

#  論理積と論理和: 2以上5以下と、11のインデックスを取得
np.where((arr4>=2) & (arr4 <=5) | (arr4 ==11))

(array([0, 1, 3, 5, 9]),)

実際に確認してみましょう。

arr4[(np.where((arr4>=2) & (arr4 <=5) | (arr4 ==11)))]

array([ 5,  2,  3, 11,  4])

確かに、「2以上5以下」の要素と「11」の要素のインデックスですね。

複数条件を1つ満たすが、同時に満たすものは除く（^演算子）

以下のコードでは、「4以下」「3である」の2つの条件を書いています。この2つを同時に満たす場合は除くので、「4以下であり3以外である数値」のインデックスを取得しています。

#  排他的論理和: 4以下で3ではない要素のインデックスを取得
np.where((arr4 <= 4) ^ (arr4 == 3))

(array([1, 2, 9]),)

実際に確認してみましょう。

arr4[np.where((arr4 <= 4) ^ (arr4 == 3))]

array([2, 1, 4])

2.3. 条件を満たす要素をx, yで置換する

オプション引数のx,yを渡すと、条件式を満たす要素をxで置換し、満たさない要素をyで置換します。

以下の配列で確認してみましょう。

import numpy as np
arr1 = np.arange(5)
print(arr1)

[0 1 2 3 4]

次のコードで、上の配列の偶数の要素を’even’に、それ以外の要素を’odd’に置換しています。

np.where(arr1%2==0, 'even', 'odd')

array(['even', 'odd', 'even', 'odd', 'even'], dtype='<U4')

2次元配列や3次元配列にも使うことができます。

arr2 = np.arange(8).reshape(2, 4)
print(arr2, '\n')

arr3 = np.where(arr2%2==0, 'even', 'odd')
print(arr3)

[[0 1 2 3]
 [4 5 6 7]] 

[['even' 'odd' 'even' 'odd']
 ['even' 'odd' 'even' 'odd']]

xとyには配列を渡すこともできます。

以下のコードだと、条件を満たす要素はそのまま、満たさない要素は-1に置換するという意味になります。

np.where(arr2%2==0, arr2, -1)

array([[ 0, -1,  2, -1],
       [ 4, -1,  6, -1]])

同時に演算を行うこともできます。

np.where(arr2%2==0, arr2*10, -1)

array([[ 0, -1, 20, -1],
       [40, -1, 60, -1]])

なお、条件に合う要素の真偽（ブール値）を確認したいだけの場合は、np.where関数を使う必要はありません。

print(arr2%2==0)

[[ True False  True False]
 [ True False  True False]]

条件が複数の場合でも同様です。

print((arr2%2==0) & (arr2>=2))

[[False False  True False]
 [ True False  True False]]

3. まとめ

以上、使いこなせるように練習しておくと良いでしょう。