如何按分组对每行计算滚动窗口内的idxmax？

问题说明 #

• 需求：按hy_code分组，获取每组最近6行数据中high列最大值对应的索引（week_date）
• 报错代码：

df=df.assign(
        l6d_highest_date=lambda x: x.groupby('hy_code')['high'].transform(lambda x: x.rolling(6).idxmax())
)

• 报错信息：AttributeError: 'Rolling' object has no attribute 'idxmax'
• 注意：数据索引为week_date

解决方案 #

Pandas的Rolling对象没有内置idxmax方法，可通过rolling().apply()自定义函数实现需求，代码如下：

def get_rolling_idxmax(series):
    # 返回滚动窗口内最大值对应的索引
    return series.idxmax()

df = df.assign(
    l6d_highest_date=lambda x: x.groupby('hy_code')['high']
                                .transform(lambda s: s.rolling(6).apply(get_rolling_idxmax, raw=False))
)

关键细节 #

• 设置raw=False，保证传入apply的是Series对象，才能调用idxmax()获取对应索引
• 若数据存在缺失值，可添加min_periods参数（如rolling(6, min_periods=1)），避免窗口数据不足时返回NaN
• 最终l6d_highest_date列会存储每组最近6行中high最大值对应的week_date索引

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