
Numpyで特定の条件に合致する値を取り出したりする際、np.whereを使うと非常に便利です。Numpyは数値計算だけでなく、PILやOpenCvなどと組み合わせることで画像処理でも使える非常に強力なライブラリで、np.whereを使いこなすことで、より高度な前処理などが可能になります。本記事ではnp.whereの使い方について解説していきます。
np.whereとは?
公式docs (https://numpy.org/doc/2.2/reference/generated/numpy.where.html)によると、以下のように書かれています。各要素と条件を比較して、基本的には条件に合致した場合、xを、それ以外はyの値を要素として格納したndarray[1]を返します。
- numpy.where(condition, [x, y, ]/)
Return elements chosen from x or y depending on condition.
(条件に応じて、xまたはyから要素を選択して返す)
[1] ndarrayはnp.arrayなどで生成される多次元配列のデータ構造のこと。
引数
np.whereの引数は、以下のように定義されています。
| 引数 | 型 | 説明 |
condition | array_like[2], bool | Trueの時xを出力し、Falseの時yを出力する |
x, y | array_like | True, Falseの時の値。どれもブロードキャスト可能なshapeである必要がある。 |
np.whereの引数[2] array_likeとはnumpyでndarrayに変換可能なデータ形式を指す。ndarrayそのものはもちろんのこと、PythonのListやTuple、スカラーもarray_likeに入る。
実際の動き
docsを見ただけではイメージをつかみにくいので、np.whereの実際の動きを見てみましょう。
conditionのみ
conditionのみを与えると、Trueになったndarrayのインデックスが返ってきます。
# Create array
In [1]: a = np.array([[0, 1, 2, 3], [4, 5, 6, 7]])
# only condition
In [2]: np.where(a > 1)
Out[2]: (array([0, 0, 1, 1, 1, 1]), array([2, 3, 0, 1, 2, 3]))
上記のように(2, 4)のndarrayを入力にすると、tupleで2つndarrayが返ってきます。それぞれのndarrayが、各軸のインデックスに対応していて、それぞれ順番に取り出してみると、
- [0, 2] -> 2
- [0. 3] -> 3
- [1, 0] -> 4
- [1, 1] -> 5
- [1, 2] -> 6
- [1, 3] -> 7
とすべて1より大きい値のインデックスが返っていることがわかります。ここで帰ってくるインデックスのサイズは、入力で使用されるndarrayのサイズに合わせて変わります。
スカラをx/yに入れる
xやyにスカラを入れると、そこで渡した値に入れ替わったndarrayが返ってきます。
# Create array
In [1]: a = np.array([[0, 1, 2, 3], [4, 5, 6, 7]])
# use scalar
In [2]: np.where(a > 1, 10, 0)
Out[2]: array([[ 0, 0, 10, 10],
[10, 10, 10, 10]])
上記のように、1よりも大きいインデックスの値が10, それ以外が0になったndarrayが返ってきます。これが割と一般的なnp.whereの使い方かなと思います。
ndarrayをx/yに入れる
ndarrayをxやyに入れることもできます。conditionで使用するndarrayと同じサイズである必要がありますが、それぞれ対応するインデックスの値に置き換えたndarrayが返ってきます。
# Create array
In [1]: a = np.array([[0, 1, 2, 3], [4, 5, 6, 7]])
In [2]: x = np.array([[10, 11, 12, 13], [14, 15, 16, 17]])
In [3]: y = np.array([[20, 21, 22, 23], [24, 25, 26, 27]])
# use array
In [4]: np.where(a > 1, in, other)
Out[4]: array([[20, 21, 12, 13],
[14, 15, 16, 17]])
上記のように、1よりも大きいインデックスの値がxのndarrayの値に置き換わり、それ以外はyのndarrayの値に置き換わっています。
ndarrayとスカラを同時に入れる
ndarrayとスカラを同時に使用することも、np.whereでは可能です。
# Create array
In [1]: a = np.array([[0, 1, 2, 3], [4, 5, 6, 7]])
In [2]: x = np.array([[10, 11, 12, 13], [14, 15, 16, 17]])
# use array and scalar
In [3]: np.where(a > 1, in, -1)
Out[4]: array([[-1, -1, 12, 13],
[14, 15, 16, 17]])
上記のように、1よりも大きいインデックスの値がxのndarrayの値に置き換わり、それ以外は-1に置き換わっています。今のndarrayの値を維持しつつ、特定の値に置換をしたい時などに便利です。
まとめ
今回はnp.whereの挙動について整理してみました。与える引数の型や、数によって、挙動が全然異なるため、使用時にはそこら辺を意識して使用できればと非常に便利なので、積極的に使ってみてください!