numpy中 where函数的用法

在我的理解中np.where函数有三个用法

1. np.where()[0] 和 np.where()[1]

where在我的理解中是一个寻找数组中某个元素的函数，在此用法中np.where()[0] 表示行索引，np.where()[1]表示列索引
具体如下

import numpy as np
array = np.arange(12).reshape(3,4)
print('array:', array)
print('np.where(array > 5):', np.where(array > 5))
print('array[np.where(a > 5)]:', array[np.where(array > 5)])
print('np.where(array > 5)[0]:', np.where(array > 5)[0])
print('np.where(array > 5)[1]:', np.where(array > 5)[1])
print(array[np.where(array > 5)[0], np.where(array > 5)[1]])

array: [[ 0 1 2 3]
 [ 4 5 6 7]
 [ 8 9 10 11]]
np.where(array > 5): (array([1, 1, 2, 2, 2, 2]), array([2, 3, 0, 1, 2, 3]))
array[np.where(array > 5)]: [ 6 7 8 9 10 11]
np.where(array > 5)[0]: [1 1 2 2 2 2]
np.where(array > 5)[1]: [2 3 0 1 2 3]
[ 6 7 8 9 10 11]

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

满足condition 这个条件的输出x，不满足输出y。
如下所示，小于2的数改为0，否则改为10

import numpy as np
t2=np.array([[1,2,3],[1,2,3]])
print(np.where(t2<2,0,10))

[[ 0 10 10]
 [ 0 10 10]]#输出

3. np.where(condition)

只有条件 (condition)，没有x和y，则输出满足条件 (即非0) 元素的坐标 (类似于numpy.nonzero)。
举个简单的例子

import numpy as np
t2=np.array([[1,2,3],[1,2,3]])
print(np.where(t2<2))

(array([0, 1], dtype=int64), array([0, 0], dtype=int64)) #输出

所以 np.where会输出每个符合要求的元素的对应的坐标

需要注意的一点是，输入的不能直接是list，需要转为array或者为array才行。比如range(10)和np.arange(10)后者返回的是数组，使用np.where才能达到效果。