如何提取DataFrame中window列的整数值传入rolling()计算可变窗口移动平均

问题根源 #

pandas的rolling()方法的window参数仅支持传入固定整数、时间偏移量等全局固定参数，不支持直接传入Series列对象为每一行指定不同的动态窗口，因此你最初的代码无法正常运行。

方案1：基于apply的易读实现（适合数据量小于10万行的场景） #

逐行读取当前行的窗口值，计算对应窗口的滚动均值：

df['sma'] = df.apply(
    lambda row: df['spx'].rolling(int(row['window'])).mean().loc[row.name],
    axis=1
)

逻辑说明：如果当前行的历史数据条数小于窗口大小，结果会自动返回NaN，符合滚动平均的常规计算规则。

方案2：基于numpy的高效实现（适合数据量大于10万行的场景） #

通过numpy数组切片直接计算，避免pandas apply的额外开销，运行效率更高：

import numpy as np

spx_values = df['spx'].values
window_values = df['window'].values
sma_result = []

for idx in range(len(df)):
    current_win = int(window_values[idx])
    # 历史数据不足窗口大小时返回NaN
    if idx + 1 < current_win:
        sma_result.append(np.nan)
    else:
        # 切片取当前行及前current_win-1行的数据计算均值
        sma_result.append(spx_values[idx - current_win + 1 : idx + 1].mean())

df['sma'] = sma_result

两种方案最终得到的sma列结果完全符合需求：比如索引15079的行自动取窗口120计算滚动平均，索引15080的行自动取窗口80计算滚动平均，和单独执行对应固定窗口rolling的结果完全一致。

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