如何高效实现DataFrame当前行与上一行不同列的取值比较计算

你当前使用逐行迭代的方案性能差的核心原因是走了Python级别的循环，对于大数据量的DataFrame，用pandas原生的向量化操作（底层C实现）可以提升10~100倍的执行效率。
另外注意原有代码存在一个隐性bug：df.sort_values(by=["start_time"]) 没有加inplace=True也没有赋值给新变量，排序结果不会生效，需要先修正。

向量化实现代码 #

import pandas as pd

def make_decision(df, max_gap):
    # 1. 按start_time排序
    df_sorted = df.sort_values(by=["start_time"]).reset_index(drop=True)
    # 2. 计算到当前行为止的最大stop_time，对应原逻辑中的tod变量
    df_sorted["cum_max_stop"] = df_sorted["stop_time"].cummax()
    # 3. 计算当前行start_time与上一行累积最大stop_time的差值，第一行差值默认设为0
    df_sorted["gap"] = df_sorted["start_time"] - df_sorted["cum_max_stop"].shift(1).fillna(df_sorted["start_time"].iloc[0])
    # 4. 差值超过max_gap的位置生成新分组标记
    df_sorted["group_id"] = (df_sorted["gap"] > max_gap).cumsum()
    # 5. 按分组聚合id列得到最终结果
    return df_sorted.groupby("group_id")["id"].apply(list).tolist()

方案说明 #

全程没有Python级别的循环，所有计算都走pandas底层优化，大体积DataFrame下性能提升非常明显，逻辑和你原有代码完全对齐：

• 按start_time排序
• 累积记录全局最大的stop_time
• 当前行start_time和上一轮的最大stop_time差值超过max_gap就拆分新分组
• 每个分组内收集对应的id列表

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