Numpy一维数组内多段不同长度累积求和的向量化实现需求

嘿，别客气！新手刚上手NumPy时碰到这类需求太正常了，完全不用道歉～我这就给你一个纯向量化、零循环的解决方案，只用NumPy自带的操作就能搞定一维数组中多段不同长度的累积求和。

核心思路 #

我们可以通过全局累积和+分段偏移量的思路实现，全程无循环：

• 先计算整个数组的全局累积和；
• 确定每个分段的起始位置，算出对应分段的“偏移值”（起始位置为0时偏移值是0，其他分段的偏移值是该分段起始前一个位置的全局累积和）；
• 把偏移值广播到对应分段的每个元素上，用全局累积和减去偏移值，就能得到每个分段独立的累积和。

代码示例与详细解释 #

下面用具体的例子来演示整个流程：

import numpy as np

# 示例输入数组
arr = np.array([1, 2, 3, 4, 5, 6, 7, 8])
# 定义各分段的长度（这里分成3段，长度依次为2、3、3）
seg_lengths = np.array([2, 3, 3])

# 1. 计算整个数组的全局累积和
full_cumsum = np.cumsum(arr)
# 输出：array([ 1,  3,  6, 10, 15, 21, 28, 36])

# 2. 计算每个分段的起始索引
starts = np.concatenate([[0], np.cumsum(seg_lengths[:-1])])
# 输出：array([0, 2, 5]) —— 对应三个分段的起始位置

# 3. 生成每个分段的偏移值，并广播到该分段的所有元素
# 第一个分段从索引0开始，偏移值为0；后续分段的偏移值是前一段最后一个元素的全局累积和
offset_values = np.concatenate([[0], full_cumsum[starts[1:] - 1]])
# 输出：array([ 0,  3, 15])
offsets = np.repeat(offset_values, seg_lengths)
# 输出：array([ 0,  0,  3,  3,  3, 15, 15, 15]) —— 每个元素对应所在分段的偏移值

# 4. 计算各分段的累积和
segmented_cumsum = full_cumsum - offsets
print(segmented_cumsum)
# 最终输出：array([ 1,  3,  3,  7, 12,  6, 13, 21])

结果验证 #

这个输出正好对应每个分段的独立累积和：

• 第一段[1,2]：累积和为 1, 1+2=3
• 第二段[3,4,5]：累积和为 3, 3+4=7, 7+5=12
• 第三段[6,7,8]：累积和为 6, 6+7=13, 13+8=21

这种方法完全依赖NumPy的向量化操作，没有任何循环，处理大规模数组时效率会非常高。

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