如何在数组中实现正负数值的分段累积求和？

可以利用NumPy的向量化操作完成需求，核心思路是先识别连续同符号的分段，再计算每段总和并将结果放置在对应分段的最后一个位置，其余位置填0。具体实现步骤如下：

步骤1：处理符号数组（含0的情况） #

先为数组每个元素确定有效符号，其中0会继承前一个最近非零元素的符号，确保连续非负/非正元素被归为同一段：

import numpy as np

# 原始数组
arr = np.array([0.2, 0.5, 1.3, 0.6, -0.3, -1.1, 0.2, -2.0, 0.7, 1.1, 0.0, -1.2])

# 初始符号数组，np.sign(0)返回0
s = np.sign(arr)
# 获取所有非零元素的索引
non_zero_idx = np.where(s != 0)[0]

# 填充0的符号：用最近的前序非零元素符号替换
if len(non_zero_idx) > 0:
    fill_idx = np.searchsorted(non_zero_idx, np.arange(len(arr)), side='right') - 1
    s = s[non_zero_idx[fill_idx]]

步骤2：识别连续分段的起止位置 #

通过符号数组的差异找到分段分界点，确定每个连续同符号段的起始和结束索引：

# 找到符号发生变化的位置
change_points = np.diff(s) != 0
# 分段起始索引（包含第一个元素）
segment_starts = np.concatenate([[0], np.where(change_points)[0] + 1])
# 分段结束索引（不包含，最后一段结束于数组长度）
segment_ends = np.concatenate([segment_starts[1:], [len(arr)]])

步骤3：计算分段总和并生成结果数组 #

用np.add.reduceat高效计算各分段总和，再将总和赋值到对应分段的最后一个位置：

# 计算每个分段的总和（无循环）
segment_sums = np.add.reduceat(arr, segment_starts)

# 初始化结果数组为全0
result = np.zeros_like(arr)
# 将分段总和赋值到对应分段的最后一个位置
result[segment_ends - 1] = segment_sums

验证结果 #

运行代码后，result的值为：

array([ 0. ,  0. ,  0. ,  2.6,  0. , -1.4,  0.2, -2. ,  0. ,  0. ,  1.8, -1.2])

与期望输出完全一致。

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