where子句与numpy

Question

这是我的数组:

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)

有人可以指导我吗？

谢谢.

Answer 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]),)

所以第二行(index == 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)

查看文档的np.allclose更多信息.