基于DataFrame指定列值设置可变窗口的Pandas滚动均值计算问题

实现方案1：逐行调用rolling（逻辑简单，适合小数据量） #

直接对每行数据读取对应窗口大小，实时计算滚动均值取当前行结果即可：

import pandas as pd
import numpy as np

# 你的原有测试代码
rng = np.random.default_rng()
df = pd.DataFrame(rng.integers(0, 100, size=(100, 2)), columns=list('AB'))
df.loc[:,'B']=df['B']//10

# 动态滚动均值计算
def get_row_rolling_mean(row, target_series):
    window_size = int(row['B'])
    if window_size < 1:
        return np.nan
    return target_series.rolling(window_size).mean().iloc[row.name]

my_series = df.apply(get_row_rolling_mean, axis=1, target_series=df['A'])

实现方案2：预计算多窗口结果（效率更高，适合大数据量） #

如果你的数据量较大，逐行计算的效率较低，可以先预计算所有出现过的窗口大小对应的滚动均值结果，再按行匹配取值：

# 提取所有去重的窗口大小
unique_windows = df['B'].unique()
window_result_map = {}

# 预计算每个窗口的滚动均值结果
for win in unique_windows:
    if win < 1:
        window_result_map[win] = pd.Series([np.nan]*len(df), index=df.index)
    else:
        # 如需窗口不足时用现有元素计算，可添加min_periods=1参数
        window_result_map[win] = df['A'].rolling(win).mean()

# 按行匹配取值
my_series = df.apply(lambda x: window_result_map[x['B']].iloc[x.name], axis=1)

注意事项 #

• 当窗口大小大于当前行的索引位置时，pandas rolling默认会返回NaN，符合常规滚动计算的逻辑
• 若你需要窗口长度不足时用已有元素计算均值，可以在rolling方法中添加min_periods=1参数
• 如果列B存在小于1的取值，你可以根据业务需求修改返回默认值，或者把窗口大小默认设置为1

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