python之带 numpy 的 Where 子句

这是我的数组:

a = [[0.,0.,0.1,0.2], [0.,0.3,0.4,0.3], [0.,0.,0.1,0.]] 

我想做一个 where 子句，它将返回 'a' 中元素的索引，其中该元素的值的总和等于 1。 类似于:where(sum(a) == 1)

有人可以指导我吗？

谢谢。

请您参考如下方法:

In [1]: import numpy as np 
 
In [2]: a = np.array([[0.,0.,0.1,0.2], [0.,0.3,0.4,0.3], [0.,0.,0.1,0.]]) 
 
In [3]: a 
Out[3]: 
array([[ 0. ,  0. ,  0.1,  0.2], 
       [ 0. ,  0.3,  0.4,  0.3], 
       [ 0. ,  0. ,  0.1,  0. ]]) 
 
In [4]: np.where(np.sum(a,axis=1) == 1) 
Out[4]: (array([1]),) 

因此第二行(索引 == 1)的总和为 1.0。 np.sum(a, axis=1) 对行求和，这相当于原始列表列表的元素。在不指定显式轴的情况下，numpy 获取数组所有元素的总和。请注意，python 内置 sum 和 np.sum 之间存在差异。这是不执行 from numpy import * 并保持明确的理由。

更新:

正如@Jaime 所建议的那样，对等式进行比较是不安全的。理想情况下 np.allclose 应该有一个 axis 选项，但它没有。您仍然可以使用以下方法重新创建它:

np.where(np.abs(np.sum(a,1) - 1.0) <= 1E-5) 

参见 docs对于 np.allclose 了解更多信息。