【NumPy】リスト内の要素で条件に合った要素のインデックスを取得したり、置換するnp.where[Python]

2024年9月18日2024年10月20日

URLをコピーしました！

NumPy

前回、PythonのNumPyで数値の正負を判別・取得するnp.sign、np.signbitを紹介しました。

条件に合った要素のインデックスを返す方法

np.whereの使い方としてまずは条件に合った要素のインデックスを返す方法を紹介します。

「np.where(ndarrayと条件)」とするとその条件に合った要素のインデックスを返します。

import numpy as np

a = np.array([-3, -2, -1, 0, 1, 2, 3])

print(np.where(a > 0))

実行結果
(array([4, 5, 6]),)

この場合、「a」のndarrayにおいて0より大きい要素のインデックスを返すことから、「4, 5, 6」が返ってくるというわけです。

ただし使えるのはndarray（NumPyの配列）のみで通常のリストではエラーとなります。

import numpy as np

b = [-3, -2, -1, 0, 1, 2, 3]

print(np.where(b > 0))

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
Cell In[3], line 5
      1 import numpy as np
      3 b = [-3, -2, -1, 0, 1, 2, 3]
----> 5 print(np.where(b > 0))

TypeError: '>' not supported between instances of 'list' and 'int'

条件によって要素を置換する方法

np.whereのもう一つの使い方として条件に合った要素を置換することもできます。

その場合は「np.where(ndarrayと条件, 条件に合った場合に置換する値, 条件に合わなかった場合に置換する値)」とします。

import numpy as np

a = np.array([-3, -2, -1, 0, 1, 2, 3])

print(np.where(a > 0, 100, 0))

実行結果
[  0   0   0   0 100 100 100]

この場合、「ndarray a」のうち「-3, -2, -1, 0」は0より大きくないため、「条件に合わなかった場合に置換する値」が使われ「0」に、「1, 2, 3」は0より大きいため、「条件に合った場合に置換する値」が使われ「100」に置換されます。

ちなみに置換したくない場合はそのリスト（ndarray）を指定します。

import numpy as np

a = np.array([-3, -2, -1, 0, 1, 2, 3])

print(np.where(a > 0, 100, a))

実行結果
[ -3  -2  -1   0 100 100 100]

import numpy as np

a = np.array([-3, -2, -1, 0, 1, 2, 3])

print(np.where(a > 0, a, -100))

実行結果
[-100 -100 -100 -100    1    2    3]

複数の条件を指定する場合

複数の条件を指定する場合、AND条件の場合は「&」、OR条件の場合は「|」を使います。

またそれぞれの条件は「()」でくくります。

import numpy as np

a = np.array([-3, -2, -1, 0, 1, 2, 3])

print(np.where((a > -2)&(a < 2), 10, -100))

実行結果
[-100 -100   10   10   10 -100 -100]

import numpy as np

a = np.array([-3, -2, -1, 0, 1, 2, 3])

print(np.where((a < -2)|(a > 2), 10, -100))

実行結果
[  10 -100 -100 -100 -100 -100   10]

ついでにNOT条件の場合は「~（チルダ）」を使います。

import numpy as np

a = np.array([-3, -2, -1, 0, 1, 2, 3])

print(np.where(a > 0, 100, 0))
print(np.where(~(a > 0), 100, 0))

実行結果
[  0   0   0   0 100 100 100]
[100 100 100 100   0   0   0]

次回はlmfitライブラリを使ってピークフィットをする方法を紹介します。