如何高效按组对DataFrame数据执行多步向前偏移操作？

嘿，你这个需求我太懂了——循环生成偏移列确实在数据量大、步数多的时候效率极低，每次对整个DataFrame赋值都会带来不必要的开销。我给你两个更高效的方案，尤其是第二个，完全用pandas的矢量化操作，速度快到飞起！

方案一：批量生成偏移列后合并（适合固定步数） #

如果你已经确定了要向前查看的固定步数（比如2步），可以一次性生成所有需要的偏移列，再合并到原数据，避免循环操作：

import pandas as pd

# 你的原始数据示例
df = pd.DataFrame({
    'user': ['Kevin', 'Kevin', 'Sara', 'Kevin'],
    'data': [1, 3, 5, 23]
})

# 定义要向前查看的步数（比如2步，对应data_1和data_2）
n_steps = 2

# 按user分组
grouped = df.groupby('user')['data']

# 用列表推导式一次性生成所有偏移列，避免循环赋值
shifted_columns = [grouped.shift(-step).rename(f'data_{step}') for step in range(1, n_steps+1)]

# 合并原数据和所有偏移列
result = pd.concat([df] + shifted_columns, axis=1)

# 保留每个user的第一条记录（匹配你想要的输出结构）
result = result.drop_duplicates(subset='user', keep='first').reset_index(drop=True)

这个方法比循环好很多，因为它只做一次concat操作，而不是多次修改整个DataFrame，减少了内存重新分配的开销。

方案二：用unstack实现完全矢量化操作（效率最高，推荐） #

这个方法利用pandas的分组编号+unstack功能，完全是底层优化的矢量化操作，不管数据量多大、步数多少，效率都拉满：

import pandas as pd

df = pd.DataFrame({
    'user': ['Kevin', 'Kevin', 'Sara', 'Kevin'],
    'data': [1, 3, 5, 23]
})

# 给每个用户的记录添加组内编号（从0开始）
df['seq'] = df.groupby('user').cumcount()

# 将组内编号作为列展开，unstack是pandas底层优化的操作，速度极快
wide_data = df.set_index(['user', 'seq'])['data'].unstack()

# 重命名列：seq=0对应原data，seq=1对应data_1，以此类推
wide_data.columns = ['data'] + [f'data_{i}' for i in range(1, len(wide_data.columns))]

# 重置索引得到最终结果
result = wide_data.reset_index()

运行后得到的结果完全符合你的需求：

如果需要固定最大步数（比如最多只保留2步），可以补充筛选列：

max_steps = 2
cols_to_keep = ['data'] + [f'data_{i}' for i in range(1, max_steps+1)]
# 自动补充不存在的列并赋值NaN
for col in cols_to_keep:
    if col not in wide_data.columns:
        wide_data[col] = pd.NA
result = wide_data[cols_to_keep].reset_index()

为什么这个方法更高效？ #

循环方法每次shift后赋值都会触发DataFrame的重新计算和内存分配，而unstack是pandas基于C语言实现的矢量化操作，直接在底层处理数据，避免了Python层面的循环开销，数据量越大，优势越明显。

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