如何在numpy中按索引累积数组?

9

我有一个数组:

a = np.array([0,0,0,0,0,0])

我希望将另一个数组添加到a的每个索引中,而该索引可以出现多次。 我想获取每个索引的总和。 我写道:

a[np.array([1,2,2,1,3])] += np.array([1,1,1,1,1])

但是要变成:

array([0, 1, 1, 1, 0, 0])

但是我想要的是:

array([0, 2, 2, 1, 0, 0])

如何在numpy中实现此操作而无需使用for循环?
4个回答

16

使用纯粹的numpy,并避免使用for循环:

np.add.at(a, np.array([1,2,2,1,3]), np.array([1,1,1,1,1]))

输出:

>>> a = np.array([0,0,0,0,0,0])
>>> np.add.at(a, np.array([1,2,2,1,3]), np.array([1,1,1,1,1]))
>>> a
array([0, 2, 2, 1, 0, 0])
请注意,这是原地替换。这是您所期望的,但未来的观众可能不希望如此。因此有此提示 :)

1
太棒了。我甚至看了np.add.reduceat,但我没有花太多时间研究ufunc方法,所以我不知道大部分方法。这里的文档甚至是完美的...“对于加法ufunc,此方法等同于a [indices] + = b,除了结果累积为被索引多次的元素” - mgilson
1
谢谢,这正是我想要的。 - maple

1
您可以自己迭代。例如:
for i in [1,2,2,1,3]:
    a[i] += 1

1
谢谢,但我想避免使用for循环。 - maple

1

我不知道有没有聪明的numpy向量化方法来做到这一点...我能想到的最好的方法是:

>>> indices = np.array([1,2,2,1,3])
>>> values = np.array([1,1,1,1,1])
>>> a = np.array([0,0,0,0,0,0])
>>> for i, ix in enumerate(indices):
...   a[ix] += values[i]
... 
>>> a
array([0, 2, 2, 1, 0, 0])

谢谢,但我想避免使用for循环。 - maple
2
@maple -- 我明白,但我不知道如何做到。当然,并不是说没有办法做到(虽然可能没有)。无论如何,有时候提供一个正确(可工作)代码示例就足以让其他人更好地解决问题。 - mgilson
@maple 如果您好奇如何解决此问题,可以检查一下我的答案。谢谢。尽管我完全同意通过举例来解决问题的做法。 - oxalorg

1
您可以这样做(假设每个索引都有对应的值):
a = np.array([0,0,0,0,0,0])
idxs = np.array([1,2,2,1,3])
vals = np.array([1,1,1,1,1])
for idx, val in zip(idxs,vals):
    a[idx] += val

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接