ndarrayの要素の値を条件に従って変更する方法

forループを使わずにndarrayの要素のうちある条件を満たすもののみを変更する方法を紹介します。大きな配列を扱う場合の処理速度はforループを使うより圧倒的に高速です。

ブーリアンマスクを使う方法

>>> a = np.array([7, 2, 4, 8, 3])
>>> a
array([7, 2, 4, 8, 3])

例として、上記のような配列に対して、 5より大きい要素を5に修正したいとします。ここでは、ブーリアンマスクを使用する方法を紹介します。

>>> mask = a > 5
>>> mask
array([ True, False, False,  True, False])

配列aを比較演算子で5と比較することによって、配列aと同じ大きさのブール型の配列が作成されました。5より大きい要素の箇所がTrueになっていることが分かります。

このブール型配列をインデックス演算子[ ]の中に入れて配列aの後ろに記述することで、5より大きい要素にのみを修正することができます。

>>> a[mask] = 5
>>> a
array([5, 2, 4, 5, 3])

なお、上記のブール型の配列maskは使うのが1回きりなので、わざわざmaskという名前を付けずに、 以下のようにまとめて書いてもOKです。

>>> a = np.array([7, 2, 4, 8, 3])
>>> a[a > 5] = 5
>>> a
array([5, 2, 4, 5, 3])

応用例1

5より大きい要素のみ2倍にする。

>>> a = np.array([7, 2, 4, 8, 3])
>>> a[a > 5] *= 2
>>> a
array([14,  2,  4, 16,  3])

応用例2

配列aの中で5より大きい要素にのみ別の配列bの値を加える。

>>> a = np.array([7, 2, 4, 8, 3])
>>> b = np.array([3, 4, 5, 6, 7])
>>> mask = a > 5
>>> a[mask] += b[mask]
>>> a
array([10,  2,  4, 14,  3])

numpy.whereを使う方法

次にnumpyの関数whereを使う方法を紹介します。

>>> a = np.array([7, 2, 4, 8, 3])
>>> np.where(a > 5, 5, a)
array([5, 2, 4, 5, 3])

上記のプログラムの2行目では、配列aのうち5を超える要素は5に、その他の要素はそのままの値とした配列を生成しています。

where関数に与えるパラメータは(condition, [x, y])となっており、conditionのところにブール型の配列を指定すると、Trueの要素はxに、Falseの要素はyになります。x,yはconditionと同じ大きさの配列か、conditionの配列にブロードキャステキングが可能なオブジェクトを指定することができます。

関連項目
配列(ndarray)の部分要素へのアクセス