如何高效实现NumPy数组的可变大小分块求和？

你可以用numpy的累加和结合索引差分的方式实现，完全避开Python循环，处理大数组时性能会提升很多：

import numpy as np

def split_sum_fast(a: np.ndarray, breakpoints: np.ndarray) -> np.ndarray:
    # 拼接完整的分割索引，包含数组首尾
    indices = np.concatenate([[0], breakpoints, [len(a)]])
    # 计算带前置0的累加和，方便直接做差分
    cum_sum = np.concatenate([[0], np.cumsum(a)])
    # 向量运算一次性算出所有块的和
    return cum_sum[indices[1:]] - cum_sum[indices[:-1]]

原理说明 #

1. 构造完整索引：把输入的断点和数组的起始（0）、结束位置（数组长度）合并，这样每个块的范围就是indices[i]到indices[i+1]，确保覆盖整个数组。
2. 累加和差分：在数组的累加和开头补一个0，这样任意区间[start, end)的元素和就等于cum_sum[end] - cum_sum[start]——这是累加和的核心性质。
3. 向量化运算：最后通过numpy的向量减法一次性计算所有块的和，全程是底层C实现的运算，没有Python循环的额外开销。

验证示例 #

a = np.array([1,2,3,4,5,6])
breakpoints = np.array([2,4])

# 原方法输出
print(split_sum(a, breakpoints))  # 输出: [ 3  7 11]
# 新方法输出
print(split_sum_fast(a, breakpoints))  # 输出: [ 3  7 11]

性能优势 #

对于百万级及以上的大数组，这个方法的速度会比原方法快几十倍甚至上百倍——原方法的Python循环需要逐个遍历每个子数组并调用np.sum，而新方法完全利用numpy的向量化特性，把所有运算都放在底层完成，大幅降低了开销。

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