python - 确定 numpy 数组的总和，同时排除某些值

Question

我想确定二维 numpy 数组的总和。但是，我想从这个总和中排除具有特定值的元素。最有效的方法是什么？

例如，这里我初始化了一个二维 numpy 数组 1 并将其中几个替换为 2:

import numpy

data_set = numpy.ones((10, 10))

data_set[4][4] = 2
data_set[5][5] = 2
data_set[6][6] = 2

如何在排除所有 2 的同时对二维数组中的元素求和？请注意，对于 10 x 10 数组，正确答案应该是 97，因为我用值 2 替换了三个元素。

我知道我可以使用嵌套的 for 循环来做到这一点。例如:

elements = []
for idx_x in range(data_set.shape[0]):
  for idx_y in range(data_set.shape[1]):
    if data_set[idx_x][idx_y] != 2:
      elements.append(data_set[idx_x][idx_y])

data_set_sum = numpy.sum(elements)

但是在我的实际数据(非常大)上，这太慢了。这样做的正确方法是什么？

Answer 1

使用 numpy 的能力 indexing with boolean arrays .在下面的示例中，data_set!=2 评估为 bool 数组，只要元素不是 2(并且具有正确的形状)，该数组就为 True。所以 data_set[data_set!=2] 是获取不包含特定值的数组的一种快速方便的方法。当然， bool 表达式可以更复杂。

In [1]: import numpy as np
In [2]: data_set = np.ones((10, 10))
In [4]: data_set[4,4] = 2
In [5]: data_set[5,5] = 2
In [6]: data_set[6,6] = 2
In [7]: data_set[data_set != 2].sum()
Out[7]: 97.0
In [8]: data_set != 2
Out[8]: 
array([[ True,  True,  True,  True,  True,  True,  True,  True,  True,
         True],
       [ True,  True,  True,  True,  True,  True,  True,  True,  True,
         True],
       ...
       [ True,  True,  True,  True,  True,  True,  True,  True,  True,
         True]], dtype=bool)