Python numpy数组对某些索引求和

Posted 2023-02-14

【中文标题】Python numpy数组对某些索引求和

【英文标题】：Python numpy array sum over certain indices

【发布时间】：2018-05-23 21:10:07

【问题描述】：

如何仅对 numpy 数组上的索引列表执行求和，例如，如果我有一个数组 a = [1,2,3,4] 和一个要求和的索引列表 indices = [0, 2] 并且我想要一个快速操作来给我答案4 因为a 中索引 0 和索引 2 处的求和值是 4

【问题讨论】：

a[indices].sum()

【参考方案1】：

试试：

>>> a = [1,2,3,4]
>>> indices = [0, 2]
>>> sum(a[i] for i in indices)
4

更快

如果你有很多数字并且想要高速，那么你需要使用numpy：

>>> import numpy as np
>>> a = np.array([1,2,3,4])
>>> a[indices]
array([1, 3])
>>> np.sum(a[indices])
4

【参考方案2】：

用indices索引后可以直接使用sum：

a = np.array([1,2,3,4])
indices = [0, 2] 
a[indices].sum()

【参考方案3】：

接受的a[indices].sum() 方法复制数据并创建一个新数组，如果数组很大，这可能会导致问题。 np.sum 实际上有一个参数来屏蔽列，你可以这样做

np.sum(a, where=[True, False, True, False])

不会复制任何数据。

掩码数组可以通过以下方式获得：

mask = np.full(4, False)
mask[np.array([0,2])] = True