Python实现滚动窗口内最大值前置最小值的计算方法

计算规则 #

• 以当前行下一行作为窗口起点，按指定长度生成滚动窗口
• 定位每个窗口内的最大值
• 若最大值位于窗口首位，直接返回该最大值
• 若最大值不在首位，返回最大值位置之前所有元素的最小值

测试用例 #

窗口长度设为3，测试序列：

import pandas as pd
import numpy as np
data = pd.Series([6,5,3,4,7,2,1])

预期输出为[5,3,4,7]，对应窗口计算逻辑：

• 窗口1[5,3,4]：最大值5在首位，返回5
• 窗口2[3,4,7]：最大值7在末位，前置元素[3,4]最小值为3，返回3
• 窗口3[4,7,2]：最大值7在中间位，前置元素[4]最小值为4，返回4
• 窗口4[7,2,1]：最大值7在首位，返回7

已有实现 #

通过numpy stride_tricks生成窗口并计算最大值的代码可正常运行，结果符合预期：

# 生成所有需要计算的窗口，跳过第一个从原数据索引0开始的窗口
windows = np.lib.stride_tricks.sliding_window_view(data.values, 3)[1:]
window_max = np.max(windows, axis=1)
# 运行结果：array([5, 7, 7, 7])

当前待解决问题为：批量提取每个窗口中最大值位置前的元素，计算对应最小值。

实现代码 #

基于已生成的窗口数组，通过向量化操作即可完成计算，无Python级循环，性能和现有stride方案一致：

# 获取每个窗口内第一个最大值的位置索引
max_pos = np.argmax(windows, axis=1)
# 生成和窗口形状一致的位置矩阵
pos_matrix = np.broadcast_to(np.arange(windows.shape[1]), windows.shape)
# 标记所有在最大值位置之前的元素
pre_mask = pos_matrix < max_pos[:, None]
# 非前置元素填充为无穷大，避免影响最小值计算
calc_arr = np.where(pre_mask, windows, np.inf)
pre_min = np.min(calc_arr, axis=1)
# 最大值在首位的位置直接取窗口最大值，其余位置取前置元素最小值
final_result = np.where(max_pos == 0, window_max, pre_min)

运行后final_result输出为array([5, 3, 4, 7])，和预期结果完全一致。

说明：np.argmax默认返回窗口内第一个出现的最大值的位置，如果业务逻辑需要取最后一个最大值的位置，替换为反向查找最大值位置的逻辑即可。

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