如何在条件触发时重启numpy的fmax累积计数？

需求：对array1执行累积最大值计算，但当当前累积最大值达到或超过同索引位置array2的元素时，将当前值重置为array1的对应元素，并从该位置重新开始累积最大值计算。

示例数据 #

import numpy as np

array2 = np.array([4,4,4,3,5,2,1])
array1 = np.array([1,2,3,2,4,1,0.5])

解决方案代码 #

# 计算原始累积最大值
orig_accum = np.fmax.accumulate(array1)
# 确定需要重置的位置：累积最大值 >= array2对应元素
reset_mask = orig_accum >= array2
# 生成拆分点，将array1划分为多个独立累积的区间
split_points = np.concatenate([[0], np.where(reset_mask)[0] + 1, [len(array1)]])
# 拆分每个区间并计算各自的累积最大值，最后拼接结果
result = np.concatenate([np.fmax.accumulate(array1[start:end]) for start, end in zip(split_points[:-1], split_points[1:])])

print(result)
# 输出：[1.  2.  3.  2.  4.  1.  0.5]

思路说明 #

1. 原始累积计算：先通过np.fmax.accumulate(array1)得到无重置的原始累积最大值数组。
2. 识别重置点：对比原始累积最大值和array2，找出需要重置的位置（累积最大值≥对应array2元素的位置）。
3. 区间拆分：以重置点为边界，将array1拆分为多个独立的子区间，每个子区间需要单独计算累积最大值。
4. 区间累积与拼接：对每个子区间执行累积最大值计算，再将结果拼接成最终数组。

这种方法完全基于Numpy的矢量化操作，没有显式的循环，效率较高。

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