Python Pandas滚动窗口最大值索引高效计算方法咨询

你的核心问题是现有基于rolling+lambda的实现效率过低，这是因为每个窗口都要执行Python层面的循环和对象操作，对于238万行的数据来说开销极大。以下是两种远快于原方案的实现方式：

方案一：Numba 加速数组运算 #

利用Numba将Python函数编译为机器码，直接操作numpy数组，避免pandas滚动窗口的Python层开销：

import numba
import numpy as np
import pandas as pd

@numba.jit(nopython=True)
def rolling_max_idx(arr, window_size):
    n = len(arr)
    result = np.full(n, np.nan)
    # 遍历每个窗口的结束位置（窗口大小达标后开始计算）
    for i in range(window_size - 1, n):
        start = i - window_size + 1
        # 截取当前窗口的数组片段
        window = arr[start:i+1]
        # 找到最大值在窗口内的位置（从1开始计数）
        max_pos = np.argmax(window) + 1
        result[i] = max_pos
    return result

# 应用到你的DataFrame，窗口大小设为1000
window_size = 1000
df['High time'] = rolling_max_idx(df['Bid'].values, window_size)

方案二：单调队列优化（O(n) 时间复杂度） #

对于超大窗口或超大数据量，单调队列可以实现线性时间复杂度的计算，每个元素仅进出队列一次，效率最高：

from collections import deque
import numpy as np

def rolling_max_idx_deque(arr, window_size):
    n = len(arr)
    result = np.full(n, np.nan)
    # 双端队列存储索引，对应arr值保持单调递减
    dq = deque()
    
    for i in range(n):
        # 移除队列中比当前元素小的索引（不可能成为后续窗口的最大值）
        while dq and arr[i] >= arr[dq[-1]]:
            dq.pop()
        # 添加当前索引到队列
        dq.append(i)
        # 移除窗口范围外的索引
        while dq[0] <= i - window_size:
            dq.popleft()
        # 窗口大小达标后，计算最大值在窗口内的位置
        if i >= window_size - 1:
            max_idx = dq[0]
            # 转换为1-1000的窗口内位置
            result[i] = i - max_idx + 1
    return result

# 应用到你的DataFrame
df['High time'] = rolling_max_idx_deque(df['Bid'].values, 1000)

原方案低效原因 #

原代码中rolling+lambda的实现，每个窗口都要创建子Series、获取索引、执行np.argmax，涉及大量Python对象的创建和销毁，属于O(n*window)的时间复杂度，在数据量较大时会产生严重的性能瓶颈。

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