moGue’s Garten

わりとアレな知識をアレな人向けになりふり構わず放出する

numpy の気づき

numpy では,こういうこともできる.

import numpy as np

A = np.array([[1, 2, -3, 4], 
              [5, -6, 7, 8], 
              [9, 10, 11, -12]])

A_sq = A ** 2
A_abs = np.abs(A)

A_positive = (A>0)
#[[ True  True False  True]
# [ True False  True  True]
# [ True  True  True False]]

X = np.empty_like(A)
X[A_positive] = A_sq[A_positive]
X[~A_positive] = A_abs[~A_positive]

print(X)
# [[  1   4   3  16]
#  [ 25   6  49  64]
#  [ 81 100 121  12]]

こういうのはスライスというのか. 配列の要素ごとに演算を行いたいが, 演算に場合分けを要する場合などに使えそうだ. ただし,A_sq および A_abs によって結局2通りの場合を全要素に対して試すことになるから,計算の無駄はある. それでも,たいていの計算ならば for 文を使うよりは速そう.