Pandas中基于groupby分组前4行及当前行生成列表列的简化实现方法

嘿，你的需求确实可以用更简洁的方式实现，尤其是amounts列的创建，完全可以用一行代码搞定，同时也能优雅地处理future列。下面是两种高效的实现方案：

方案1：使用Rolling窗口（直观易读） #

import pandas as pd

data_test = {'nr':[1,1,1,1,1,6,6,6,6,6,6,6], 'val':[11,12,13,14,15,61,62,63,64,65,66,67]}
df_test = pd.DataFrame(data_test, columns=['nr','val'])

# 一行代码生成amounts列
df_test['amounts'] = df_test.groupby('nr')['val'].rolling(window=5, min_periods=1).apply(
    lambda x: [0]*(5 - len(x)) + x.tolist(), raw=False
).reset_index(level=0, drop=True)

# 生成future列（和你原来的逻辑一致）
df_test['future'] = df_test.groupby('nr')['val'].shift(-1).fillna(0)

print(df_test)

代码解释： #

• rolling(window=5)：对每个分组的val字段设置滑动窗口，窗口大小为5，刚好覆盖当前行+前4行
• min_periods=1：允许窗口中最少包含1个值，保证分组开头的行（比如前4行）也能正常生成列表
• apply逻辑：如果窗口长度不足5，就在列表前面补0，再拼接窗口内的实际值，和你原来的amounts列输出完全一致
• reset_index(level=0, drop=True)：移除分组索引，让结果和原DataFrame的行索引对齐

方案2：使用Shift列表推导（性能更优） #

如果你的数据集很大，rolling+apply的性能可能不够理想，可以用这种基于shift的一行实现，性能更高效：

df_test['amounts'] = df_test.groupby('nr')['val'].apply(
    lambda x: pd.concat([x.shift(i) for i in range(4, -1, -1)], axis=1).fillna(0).values.tolist()
).explode()

代码解释： #

• 用列表推导式生成shift(4)到shift(0)（即当前值）的5列数据
• 拼接这5列后填充0，再转成列表格式
• explode()把每个分组对应的列表展开，和原DataFrame的每行一一对应

两种方案生成的amounts和future列结果和你原来的实现完全一致，但代码更简洁、易维护，不需要手动创建多个中间列。

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