按索引对numpy数组的累积求和

Posted 2023-02-14

【中文标题】按索引对numpy数组的累积求和

【英文标题】：Cumulative summation of a numpy array by index

【发布时间】：2011-04-06 02:00:17

【问题描述】：

假设您有一个需要相加的值数组

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

第二个数组指定哪些元素需要相加

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

结果将存储在大小为max(i)+1 的新数组中。因此，例如i=[0,0,0,0,0] 相当于将d 的所有元素相加并将结果存储在大小为1 的新数组的0 位置。

我尝试使用

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

但是，+= 操作仅将每个元素添加一次，从而给出了意外的结果

[1,1,1]

而不是

[2,1,2]

如何正确实现这种求和？

【问题讨论】：

如果d 的值是唯一的，这会更清楚。例如，如果d = [0,1,2,3,4] Im guessing for i = [0,0,0,0,0]` 你想要c = [10]，而对于i = [0,0,1,2,2] 你想要c = [1,2,7]？

没错。感谢您的澄清。

在这种情况下，juxstapose 的解决方案以及我在 cmets 中建议的更改应该可以解决问题。

【参考方案1】：

如果我理解正确的话，有一个快速的函数（只要数据数组是 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 返回所有整数 range(max(i)) 的数组，即使某些计数为零

【讨论】：

这是这里描述的案例的最佳解决方案。对于标记数组的一般总和，您可以使用 scipy.ndimage.sum。该模块还具有其他有用的功能，例如最大值、最小值、均值、方差、...

【参考方案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))

【参考方案3】：

此解决方案对于大型数组应该更有效（它迭代可能的索引值而不是 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.]

【参考方案4】：

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

【讨论】：

关闭，但我认为 OP 想要for didx,ridx in enumerate(i_list): result[ridx] += d[didx]。此外，由于标签包含 [numpy]，您可以使用 numpy.zeros。

【参考方案5】：

在一般情况下，当您想按标签对子矩阵求和时，您可以使用以下代码

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