NumPy数组按多组索引批量求和的高效向量化实现方法

输入定义 #

• 待操作的2D NumPy数组：

import numpy as np
m = np.array([[0,0,0,0], [0, 0.08, 0.1, 0.02], [0,0.12,0.15,0.03], [0,0.2,0.25,0.05]])

• 分组索引列表，每个元素对应一组需要提取求和的坐标：

filter_indices=[np.array([[1,1],[3,1],[1,3]]), np.array([[2,2],[2,3],[2,3]]), np.array([3,3]) ]

• 预期输出：一维数组array([0.3 , 0.21, 0.05])，要求避免显式Python层循环遍历分组做索引求和，提升运算效率。

实现方案 #

采用统一索引+分组聚合的思路，核心计算全部在NumPy C层完成，兼容不同长度的索引分组、单坐标一维数组的写法：

# 为每个坐标标记所属分组ID，拼接所有坐标为统一的二维数组
group_ids = np.concatenate([
    np.full(len(np.atleast_2d(ind)), idx, dtype=np.intp) 
    for idx, ind in enumerate(filter_indices)
])
all_coords = np.concatenate([np.atleast_2d(ind) for ind in filter_indices])
# 一次性提取所有目标位置的值，按分组ID加权求和
result = np.bincount(group_ids, weights=m[all_coords[:, 0], all_coords[:, 1]])

效果说明 #

• 运行代码得到的result与原列表推导式输出完全一致，为array([0.3 , 0.21, 0.05])
• 该实现仅在生成分组ID、拼接坐标时有极轻量的构造开销，避免了原方案中循环内重复触发NumPy索引、求和操作的Python层开销，分组数量越大性能优势越明显
• 如果所有索引组的长度完全相等，可以直接将索引列表转为shape为(组数, 每组索引数, 2)的三维数组，批量索引后沿对应轴求和，性能还可进一步提升，上述方案为适配不等长分组的通用最优解。

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