np.where用法

官方文档
两种用法

1.确定一个array中满足条件的索引列表

>>> import numpy as np
>>> a = np.arange(10)
>>> np.where(a < 5)
(array([0, 1, 2, 3, 4], dtype=int64),)

注意这里返回的是一个元组，需要获取到索引列表还需要[0]来访问返回结果的第一个元素

2.根据条件返回和原array形状一致的array

>>> import numpy as np
>>> a = np.arange(10)
>>> np.where(a < 5, a, 10*a)
array([ 0,  1,  2,  3,  4, 50, 60, 70, 80, 90])

三个参数：条件，条件为真的取值，条件为假的取值
直接返回一个array