python - 带有 numpy 的 Where 子句

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

所以第 2 行（索引 == 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

Answer 2

使用enumerate，列表理解：

>>> 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]