numpy模块中的where函数

numpy.where — NumPy v1.20 Manual

numpy.where(condition, x, y)

如果条件成立，执行x；如果条件不成立，执行y

例1：

import numpy 
a = numpy.array([1, 2, 3, 4, 5, 6, 7])
numpy.where(a>5, 0, 1)

返回的结果是array([1, 1, 1, 1, 1, 0, 0])

例2：

import numpy 
a = numpy.array([1, 2, 3, 4, 5, 6, 7])
numpy.where(a>5)

返回的结果是(array([5, 6], dtype=int64),)

注：这里的5和6是值所对应的索引

例3：

import numpy 
 np.where([[True, False], [True, True]],
           [[1, 2], [3, 4]],
           [[9, 8], [7, 6]])

返回的结果是array([[1, 8],[3, 4]])

注：返回的结果与condition这里的shape相同，可以这样解读：值为True则取x中的元素，值为False取y中的元素