无需Lambda函数：Pandas按ID分组计算前三行差值的高效实现

给定如下示例DataFrame，实际数据包含数百万行，已按timestamp排序但ID列混乱：

import pandas as pd

data = {
    'ID': [1, 1, 2, 1, 2, 1, 1, 2, 2, 2, 1, 2, 2, 1],
    'timestamp': pd.date_range(start='1/1/2023', periods=14, freq='D'),
    'value': [11, 22, 33, 44, 55, 66, 77, 88, 99, 11, 22, 33, 44, 55]
}
df = pd.DataFrame(data)

需求是按ID分组，计算每行与组内前三行的value差值，最终根据差值判断增减并过滤数据。原实现用groupby+apply(lambda)效率较低，以下是更高效的优化方案：

优化方案：利用groupby+shift的向量化操作 #

pandas的groupby结合shift支持直接在组内进行偏移计算，无需借助lambda函数，完全是向量化操作，能大幅提升百万级数据的处理效率。

完整步骤 #

1. 先按ID和timestamp排序，确保组内数据按时间顺序排列：

df = df.sort_values(by=['ID', 'timestamp'])

2. 直接计算组内前三行的value差值：

# 计算当前value与组内前三行value的差值
df['value_diff'] = df['value'] - df.groupby('ID')['value'].shift(3)

3. 根据差值过滤数据（例如保留差值非负的行，组内前三行的value_diff为NaN，可按需处理）：

filtered_df = df[df['value_diff'].ge(0) | df['value_diff'].isna()]

效率优势说明 #

原方案的apply(lambda x: x - x.shift(3))本质是逐组调用Python函数，会产生额外的函数调用开销；而优化后的groupby+shift是pandas底层的向量化实现，直接对整列进行操作，避免了逐组循环，在百万级数据场景下速度能提升数倍甚至一个量级。

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