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 {{1}}。

Answer 2

使用enumerate，list comprehension：

>>> a = [[0.,0.,0.1,0.2], [0.,0.3,0.4,0.3], [0.,0.,0.1,0.]]
>>> [i for i, xs in enumerate(a) if sum(xs) == 1]
[1]

where子句与numpy

2 个答案: