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]