Numpy中按索引数组分组求和的操作名称与实现方法

这种操作叫做分组求和（也可称为“按索引分组的归约求和”），属于归约操作的一种，即根据给定的分组索引将数组元素分组后，对每组内的元素进行求和聚合。

以下是两种高效的实现方式：

方法一：使用 np.bincount #

这是最直接的方法，适用于分组索引（数组 b）为从 0 开始的非负整数的场景（正好符合你的需求，b 的元素范围是 [0, M-1]）。np.bincount 可以接收权重参数，直接对每组对应的元素求和。

示例代码：

import numpy as np

a = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
b = np.array([3, 3, 0, 3, 2, 1, 2, 1, 2, 2])
c = np.bincount(b, weights=a)

print(c)  # 输出: [ 2. 12. 27.  4.]

方法二：使用 np.add.at #

这种方法更灵活，适合需要自定义分组操作的场景。通过初始化一个全零数组，再利用 np.add.at 对指定索引位置进行累加操作。

示例代码：

import numpy as np

a = np.array([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
b = np.array([3, 3, 0, 3, 2, 1, 2, 1, 2, 2])
M = 4  # 分组总数
c = np.zeros(M, dtype=np.float64)

np.add.at(c, b, a)
print(c)  # 输出: [ 2. 12. 27.  4.]

内容的提问来源于stack exchange，提问作者RRR