Python中如何使用lag函数分组位移数据生成多列滞后变量

首先确保已安装pandas库，核心实现逻辑如下：

import pandas as pd

# 构建测试数据集
data = [
    [1,1,38147],
    [2,1,567],
    [3,1,0],
    [4,1,0],
    [5,1,0],
    [6,1,0],
    [7,1,0],
    [8,1,0],
    [1,2,3099],
    [2,2,30460],
    [3,2,2372],
    [4,2,0],
    [5,2,0],
    [6,2,0],
    [7,2,0],
    [8,2,267],
]
test1 = pd.DataFrame(data, columns = ['FIRST', 'SECOND', 'SCORE'])

# 动态获取需要生成的滞后变量数量，取FIRST列的唯一值数量即可适配动态调整的场景
num_lags = test1['FIRST'].nunique()

# 按SECOND分组，依次生成对应阶数的滞后列
for lag in range(1, num_lags + 1):
    test1[f'S{lag}'] = test1.groupby('SECOND')['SCORE'].shift(lag)

# 打印验证输出结果
print(test1)

逻辑说明 #

• 没有硬编码滞后阶数，通过nunique()动态获取取值数量，FIRST/SECOND的唯一值数量变化时不需要修改核心代码
• groupby('SECOND')限制滞后计算仅在同一个SECOND分组内执行，和参考R代码的by = SECOND逻辑完全一致
• shift(lag)原生默认用NaN填充缺失值，输出结果和需求完全匹配

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