按索引计算的numpy数组的累积总和

时间:2010-08-31 04:32:13

标签: python numpy sum indices

假设您有一组需要一起求和的值

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 个答案:

答案 0 :(得分:11)

如果我正确理解了这个问题,那么就有一个快速的功能(只要数据数组为1d)

>>> 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))的数组,即使某些计数为零

答案 1 :(得分:2)

def zeros(ilen):
 r = []
 for i in range(0,ilen):
     r.append(0)

i_list = [0,0,1,2,2]
d = [1,1,1,1,1]
result = zeros(max(i_list)+1)

for index in i_list:
  result[index]+=d[index]

print result

答案 2 :(得分: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.]

答案 3 :(得分:2)

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

答案 4 :(得分:0)

在一般情况下,如果要按标签对子矩阵求和,可以使用以下代码

import numpy as np
from scipy.sparse import coo_matrix

def labeled_sum1(x, labels):
     P = coo_matrix((np.ones(x.shape[0]), (labels, np.arange(len(labels)))))
     res = P.dot(x.reshape((x.shape[0], np.prod(x.shape[1:]))))
     return res.reshape((res.shape[0],) + x.shape[1:])

def labeled_sum2(x, labels):
     res = np.empty((np.max(labels) + 1,) + x.shape[1:], x.dtype)
     for i in np.ndindex(x.shape[1:]):
         res[(...,)+i] = np.bincount(labels, x[(...,)+i])
     return res

第一种方法使用稀疏矩阵乘法。第二个是user333700答案的概括。两种方法的速度都相当:

x = np.random.randn(100000, 10, 10)
labels = np.random.randint(0, 1000, 100000)
%time res1 = labeled_sum1(x, labels)
%time res2 = labeled_sum2(x, labels)
np.all(res1 == res2)

输出:

Wall time: 73.2 ms
Wall time: 68.9 ms
True