NumPy配列での三項演算子ベクトル計算版 2015/05/28

np.where は三項演算子のベクトル計算といえる。

3つの引数を取り、1つ目が条件式、つまり論理値のndarray、2つ目が条件がtrueのときの値、3つ目が条件がfalseのときの値。

例その1

ndarr1 = np.array([1, 2, 3, 4, 5])
ndarr2 = np.array([10, 20, 30, 40, 50])
cond = np.array([True, False, True, True, False])
print(np.where(cond, ndarr1, ndarr2))
# 出力結果
# [ 1 20  3  4 50]

例その2

ndarr1 = np.random.rand(4, 2) - 0.5
ndarr2 = np.zeros((4, 2), dtype='float64')

print(ndarr1)
# 出力結果
# [[-0.32185056  0.0451717 ]
#  [ 0.04564159 -0.15487149]
#  [-0.04431453 -0.45999036]
#  [-0.3508403   0.0938512 ]]

print(ndarr2)
# 出力結果
# [[ 0.  0.]
#  [ 0.  0.]
#  [ 0.  0.]
#  [ 0.  0.]]

ndarr3 = np.where(ndarr1 > 0.0, ndarr1, ndarr2)

print(ndarr3)
# 出力結果
# [[ 0.          0.0451717 ]
#  [ 0.04564159  0.        ]
#  [ 0.          0.        ]
#  [ 0.          0.0938512 ]]

np.random.rand(4, 2) - 0.5 は 4 × 2 の行列の各要素を0から1までの乱数にして、さらにベクトル計算のブロードキャストにより各要素から 0.5 を引いた行列を生成する。

次の例のように2つ目や3つ目の引数はndarrayでなくてスカラーでも構わない。これは np.where でベクトル計算のブロードキャストするイメージ。

例その3

ndarr1 = np.random.rand(4, 2) - 0.5

print(ndarr1)
# 出力結果
# [[-0.48128308  0.27579344]
#  [ 0.36287708  0.26540085]
#  [ 0.42619999  0.29142238]
#  [-0.15820554 -0.35493536]]

ndarr3 = np.where(ndarr1 > 0.0, ndarr1, 0.0)

print(ndarr3)
# 出力結果
# [[ 0.          0.27579344]
#  [ 0.36287708  0.26540085]
#  [ 0.42619999  0.29142238]
#  [ 0.          0.        ]]