Pandas按客户分组计算不同窗口大小的滚动特征

给定原始客户周期数据：

import pandas as pd

df = pd.DataFrame({'cust_id': [1,1,1,1,1,1,2,2,2,2,2,2],
                   'period' : [1,2,3,4,5,6,1,2,3,4,5,6],
                   'volume' : [1,2,3,4,5,6,7,8,9,10,11,12]})

需要按客户分别计算窗口大小为3和6的滚动最大值，期望输出如下：

out = pd.DataFrame({'cust_id': [1,1,1,1,1,1,2,2,2,2,2,2],
                    'max_vol_3' : [None, None, 3,4,5,6,None,None,9,10,11,12],
                    'max_vol_6' : [None,None,None,None,None,6,None,None,None,None,None,12]})

解决方案 #

不需要透视操作，直接通过groupby结合rolling就能实现，代码如下：

import pandas as pd

# 初始化原始数据
df = pd.DataFrame({'cust_id': [1,1,1,1,1,1,2,2,2,2,2,2],
                   'period' : [1,2,3,4,5,6,1,2,3,4,5,6],
                   'volume' : [1,2,3,4,5,6,7,8,9,10,11,12]})

# 计算窗口3的滚动最大值（按客户分组）
df['max_vol_3'] = df.groupby('cust_id')['volume'].rolling(window=3, min_periods=3).max().reset_index(level=0, drop=True)

# 计算窗口6的滚动最大值（按客户分组）
df['max_vol_6'] = df.groupby('cust_id')['volume'].rolling(window=6, min_periods=6).max().reset_index(level=0, drop=True)

# 提取目标列得到最终结果
out = df[['cust_id', 'max_vol_3', 'max_vol_6']]

代码说明 #

• groupby('cust_id')：确保每个客户的滚动计算独立进行，避免跨客户的数据干扰
• rolling(window=N, min_periods=N)：设置滚动窗口大小为N，同时要求窗口内必须有N个有效数据才计算最大值，不足时返回NaN（与期望输出中的None等价）
• reset_index(level=0, drop=True)：移除分组产生的额外索引层级，让计算结果与原DataFrame的行完美对齐

运行上述代码后，得到的out与期望输出完全一致。

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