np.where()

numpy.where (condition[, x, y])

 

numpy.where()兩種用法

1. np.where(condition, x, y)

滿足條件(condition)，輸出x，不滿足輸出y。

import numpy as np

a = np.arange(10)
b = np.where(a>5, 1, -1)
c = np.where(a, 1, -1)    # 0為False，所以第一個輸出-1
print(b)
print(c)

[-1 -1 -1 -1 -1 -1  1  1  1  1]
[-1  1  1  1  1  1  1  1  1  1]

 

2. np.where(condition)

只有條件 (condition)，沒有x和y，則輸出滿足條件 (即非0) 元素的坐標.

a = np.arange(0, 100, 10)
b = np.where(a < 50) 
c = np.where(a >= 50)[0]
print(b) 
print(c)

(array([0, 1, 2, 3, 4], dtype=int64),)
[5 6 7 8 9]