Numpy数组按索引分组求和内存占用过高，应如何优化实现？

你当前代码运行慢、内存占用超限的核心原因是列表推导式本质为Python层循环，800-900万次循环的解释器开销极高，且每次循环生成的A_numpy_array[b]临时数组无法及时释放，累积占用直接触发内存溢出。可根据B数组的结构选择以下两种完全向量化的实现，内存效率和运行速度都会有数量级提升：

直接用numpy高级索引一次性取值，沿对应维度求和即可，全程无Python层循环：

# 每个b对应A的第一维索引，求和后长度和B的第一维一致
df[col] = A_numpy_array[B_numpy_array].sum(axis=1)

该方式内存开销仅为最终结果的1-2倍，运行速度是原列表推导式的百倍以上。

用np.add.reduceat做分段求和，仅需要一次轻量的长度统计循环，其余逻辑全在numpy层执行：

# 统计每个索引组的长度
elem_lens = np.array([len(b) for b in B_numpy_array])
# 拼接所有索引为一维数组
all_indices = np.concatenate(B_numpy_array)
# 计算分段求和的起始位置
split_points = np.r_[0, elem_lens.cumsum()[:-1]]
# 一次性取所有对应值后分段求和
df[col] = np.add.reduceat(A_numpy_array[all_indices], split_points)

如果A是多维数组，原逻辑需要对每个b取出的子数组全量求和，可修改取值逻辑为：

all_vals = A_numpy_array[all_indices].reshape(len(all_indices), -1).sum(axis=1)
df[col] = np.add.reduceat(all_vals, split_points)

该方式避免了每次循环生成临时数组和重复做求和运算，内存峰值仅为所有索引值+所有对应A值的总大小，远低于原实现的内存占用。

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