https://docs.scipy.org/doc/numpy/reference/generated/numpy.where.html
numpy.where(condition,x,y)
对于条件condition,满足则返回x,不满足返回y。
示例:
print(np.where([[True, False], [True, True]],[[1, 2], [3, 4]],[[9, 8], [7, 6]]))
上面代码返回值为[[1,8],[3,4]],[True, False]中第一个元素True满足条件,x的第一个元素为1,第二个元素为False,y中第二个元素为8,同样地,[True, True]分别对应x中的[3,4].
如果x,y为缺省值,则返回满足条件的下标。
x = np.arange(9.).reshape(3, 3)
print(np.where( x > 4 ))
print(np.where( x > 5 ))
print(np.where( x > 6 ))
上面的三个输出分别为
(array([1, 2, 2, 2]), array([2, 0, 1, 2])) ,
(array([2, 2, 2]), array([0, 1, 2]))
(array([2, 2]), array([1, 2]))
对于np.where( x > 4 ),大于4的值为5,6,7,8,下标分别为(1,2),(2,0),(2,1),(2,2),第一维度为(1,2,2,2),第二维度为(2,0,1,2).