按组计算Pandas DataFrame行差值：递增周期至预设值的实现

给定如下DataFrame：

import pandas as pd
df = pd.DataFrame([['a', 3], ['a', 5], ['a', 2], ['a', 6], ['a', 7], ['a', 1], ['a', 9], ['b', 7], ['b', 8], ['b', 11], ['b', 9], ['b', 10], ['b', 6]], columns = ['k', 'v'])

需求：

• 按k列分组后，对每组的v列计算行差值
• 预设周期为3，即从第4行开始，取当前行与前3行的差值
• 每组前3行不足3个前置值时，偏移量从1递增至3：第1行无前置值，差值为NaN；第2行取与前1行的差值；第3行取与前2行的差值

基础实现版本 #

通过分组内索引标记、偏移量计算、分组差值计算三步完成，完全遵循Pandas向量化操作逻辑：

import pandas as pd

# 1. 为每组生成从0开始的行索引
df['row_idx'] = df.groupby('k').cumcount()

# 2. 计算每行对应的偏移量：前3行偏移量为行号+1，之后固定为3
df['offset'] = df['row_idx'].apply(lambda x: x + 1 if x + 1 <= 3 else 3)

# 3. 分组后按偏移量计算v列的差值：当前值减去对应偏移量前的值
df['diff'] = df.groupby('k').apply(
    lambda group: group['v'] - group['v'].shift(group['offset'])
).reset_index(level=0, drop=True)

# 清理辅助列（可选）
df = df.drop(['row_idx', 'offset'], axis=1)

执行后结果如下：

高效向量化优化版本 #

针对大数据集，可替换为更高效的向量化操作，避免apply中的lambda循环：

import pandas as pd
import numpy as np

df['row_idx'] = df.groupby('k').cumcount()
# 用numpy.where替代apply计算偏移量，性能更优
df['offset'] = np.where(df['row_idx'] < 3, df['row_idx'] + 1, 3)

# 用groupby.transform实现分组内的偏移差值计算，保持索引对齐
df['diff'] = df['v'] - df.groupby('k')['v'].transform(
    lambda s: s.shift(s.reset_index(drop=True).index.to_series().apply(lambda x: x+1 if x+1<=3 else 3))
)

# 清理辅助列
df = df.drop(['row_idx', 'offset'], axis=1)

这个版本利用transform直接返回与原DataFrame对齐的结果，避免了索引重置操作，在数据量较大时性能提升明显。

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