基于索引数组的NumPy数组求和（或均值）实现问询

针对大数组的按索引分组统计需求，推荐用NumPy的bincount函数实现，完全避免Python循环，性能拉满。

核心实现（求和） #

import numpy as np

value_array = np.array([[0.1, 0.2, 0.5, 1 ],
                       [2  ,   5,  10, 20]])
index_array = np.array([[  0,   1,   1, 2 ],
                       [  2,   0,   1, 4 ]])

# 展平二维数组为一维，适配bincount的输入要求
flat_values = value_array.flatten()
flat_indices = index_array.flatten()

# 按索引分组求和，minlength确保覆盖所有出现过的索引（包括最大索引）
desired_sum = np.bincount(flat_indices, weights=flat_values, minlength=np.max(flat_indices) + 1)

print(desired_sum)
# 输出：array([5.1, 10.7, 3. , 0. , 20.])

计算均值的扩展 #

如果需要分组均值，先统计每个索引的元素个数，再用总和除以个数：

# 统计每个索引的元素数量
counts = np.bincount(flat_indices, minlength=np.max(flat_indices) + 1)
# 处理除数为0的情况（无元素的索引均值设为0）
desired_mean = np.where(counts != 0, desired_sum / counts, 0)

print(desired_mean)
# 输出：array([2.55, 3.56666667, 1.5, 0., 20.])

性能说明 #

bincount是NumPy底层优化的C实现，处理(2000,2000)规模的数组（400万元素）毫无压力，比Python循环快数十倍甚至上百倍，完全满足你的性能需求。

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