np.where说明

引用几个例子

>>> a = np.array([-1, -2, 0, 1, 2])
>>> np.where(a)
(array([0, 1, 3, 4]),)

只有input，后面没有指定条件，返回非0的元素

>>>a = np.arange(10)
>>>a
array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>>np.where(a < 5, a, 10*a)
array([ 0,  1,  2,  3,  4, 50, 60, 70, 80, 90])

有指定条件，等价为 if a的元素 < 5, 元素 =它本身, else 元素 = 10*它本身

>>> np.where([[True, False], [True, True]],
         [[1, 2], [3, 4]],
         [[9, 8], [7, 6]])
output:         
array([[1, 8],
       [3, 4]])

这里相当于一个2x2数组(T, F ; T, T), 其中" ; "表示换行，它后面跟的有两个2x2数组，分别对应了if T 条件 和else的条件。
具体来说，
如果第一行的(T, F)的元素为(T, T)，它就等于(1, 2), 否则等于(9, 8)
现在第一个元素为T，就等于1, 第2个元素为F，就等于8。所以第一行返回(1, 8)
同理第二行。

再举个例子，

>>> b = np.array([[1,1],[2,3] ,[0.2,0.5]])
>>> b
array([[1. , 1. ],
       [2. , 3. ],
       [0.2, 0.5]])

>>> np.where(b)
output:
(array([0, 0, 1, 1, 2, 2]), array([0, 1, 0, 1, 0, 1]))

你肯定想说，为什么返回两个一维数组，这是什么意思。
首先，b是一个3x2数组，np.where要返回b中> 0的元素位置
位置包含了行坐标和列坐标，
那第一行就是行坐标，第2行就是列坐标，不信把它们拼在一起看一下

0, 0, 1, 1, 2, 2
0, 1, 0, 1, 0, 1

按列看，是不是坐标就是(0,0), (0, 1), (1, 0), … , (2,1)

把b改一下，把最后一个元素改为0

>>> b
array([[1. , 1. ],
       [2. , 3. ],
       [0.2, 0. ]])

>>> np.where(b)
output:
(array([0, 0, 1, 1, 2]), array([0, 1, 0, 1, 0]))

可以看到，最后一个元素的坐标没有了。