按索引对numpy数组进行累计求和

12 投票
5 回答
8299 浏览
提问于 2025-04-16 03:30

假设你有一个数值数组,需要把这些数值加在一起。

d = [1,1,1,1,1]

还有一个第二个数组,里面指定了哪些元素需要被加在一起。

i = [0,0,1,2,2]

结果会存储在一个新的数组里,这个数组的大小是 max(i)+1。举个例子,如果 i=[0,0,0,0,0],那么就相当于把所有 d 中的元素加起来,然后把结果放在新数组的第 0 个位置,这个新数组的大小是 1

我尝试用下面的方式来实现这个功能:

c = zeros(max(i)+1)
c[i] += d

但是,+= 这个操作只会把每个元素加一次,所以得到了意想不到的结果:

[1,1,1]

而不是:

[2,1,2]

那要怎么正确地实现这种加法呢?

5 个回答

3

Juh_的评论是最有效的解决方案。下面是可以运行的代码:

import numpy as np
import scipy.ndimage as ni

i = np.array([0,0,1,2,2])
d = np.array([0,1,2,3,4])

n_indices = i.max() + 1
print ni.sum(d, i, np.arange(n_indices))
14

如果我理解这个问题没错的话,这里有一个很快的函数可以用(只要数据是一个一维的数组)

>>> i = np.array([0,0,1,2,2])
>>> d = np.array([0,1,2,3,4])
>>> np.bincount(i, weights=d)
array([ 1.,  2.,  7.])

np.bincount 会返回一个数组,包含所有整数的范围(最大值是 max(i)),即使有些计数是零也会显示出来

2

这个解决方案对于处理大数组来说应该更有效率,因为它是遍历可能的索引值,而不是逐个检查每个i的条目。

import numpy as np

i = np.array([0,0,1,2,2])
d = np.array([0,1,2,3,4])

i_max = i.max()
c = np.empty(i_max+1)
for j in range(i_max+1):
    c[j] = d[i==j].sum()

print c
[1. 2. 7.]

撰写回答