如何用Numpy替代循环实现带条件重置的数组累积最大值计算

问题需求 #

对数组longband计算累积最大值，但当**不满足longband[i-1] > longband[i]且longband[i-1] < RsiMa[i]**时，需保留longband[i]的原始值并重新开始累积。目前仅通过Python for循环实现，希望改用Numpy实现并了解具体方法。

原循环实现代码 #

import numpy as np
RsiMa = np.array([2, 4, 6, 9, 10, 8, 6, 4, 6, 8, 10, 12, 10, 8])
longband = np.array([1, 3, 2, 8, 7, 5, 4, 3, 4, 3, 7, 10, 7, 6])

for i in range(1, len(longband)):
    if i > 1:
        if longband[i-1] > longband[i] and longband[i-1] < RsiMa[i]:
            longband[i] = longband[i-1]

预期输出 #

[ 1  3  3  8  8  5  5  3  4  4  7 10  7  7 ]

已尝试方案 #

• np.fmax.accumulate()：无法根据指定条件触发累积重置，不符合需求。
• Numba JIT优化：运行效率不升反降，耗时从15ms增至33ms。

Numpy实现方案 #

核心思路是标记累积分组：先确定需要重置累积的位置，将数组划分为多个连续的累积段，再对每个段单独计算累积最大值，具体步骤如下：

import numpy as np

RsiMa = np.array([2, 4, 6, 9, 10, 8, 6, 4, 6, 8, 10, 12, 10, 8])
longband = np.array([1, 3, 2, 8, 7, 5, 4, 3, 4, 3, 7, 10, 7, 6])

# 1. 标记重置点：不满足累积条件时，开启新的累积分组
reset_mask = np.zeros(len(longband), dtype=bool)
reset_mask[0] = True  # 第一个元素作为初始分组起点
reset_mask[1] = False # 原循环中i=1时不执行替换，属于初始分组

for i in range(2, len(longband)):
    prev_val = longband[i-1]
    curr_val = longband[i]
    # 不满足累积条件则标记为重置点
    if not (prev_val > curr_val and prev_val < RsiMa[i]):
        reset_mask[i] = True

# 2. 生成分组ID：每个连续非重置段属于同一组
group_ids = np.cumsum(reset_mask)

# 3. 对每个分组单独计算累积最大值
result = np.zeros_like(longband)
for group in np.unique(group_ids):
    segment_mask = group_ids == group
    result[segment_mask] = np.maximum.accumulate(longband[segment_mask])

print(result)

输出验证 #

运行代码后输出与预期完全一致：

[ 1  3  3  8  8  5  5  3  4  4  7 10  7  7 ]

效率说明 #

该方案利用Numpy向量化操作处理分组内的累积计算，仅在分组数量较少时存在少量循环开销，相比纯Python循环，在百万级以上规模的数组上会有明显速度提升。

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