Python pandas按两列分组且行时间差小于2秒时合并行的实现问题

import pandas as pd
import datetime

def merger(dataframe: pd.DataFrame) -> pd.DataFrame:
    # 预处理：复制原表避免修改原数据，转换时间格式，计算秒级时间戳
    df = dataframe.copy()
    df['DATE_TIME'] = pd.to_datetime(df['DATE_TIME'].str.strip("'"))
    df['epoch'] = df['DATE_TIME'].astype('int64') // 10**9

    # 按通话双方分组，组内按通话开始时间升序排序
    df = df.sort_values(by=['A_PERSON', 'B_PERSON', 'DATE_TIME']).reset_index(drop=True)

    # 同组内判断是否需要新起合并分组：当前行开始时间与上一行结束时间差值≥2秒则新分组
    df['prev_end'] = df.groupby(['A_PERSON', 'B_PERSON'])['epoch'].shift(1) + df.groupby(['A_PERSON', 'B_PERSON'])['DURATION'].shift(1)
    df['new_group'] = (df['epoch'] - df['prev_end'] >= 2).fillna(True)
    df['group_id'] = df.groupby(['A_PERSON', 'B_PERSON'])['new_group'].cumsum()

    # 按分组聚合：开始时间取最早值，通话时长取总和
    res_df = df.groupby(['A_PERSON', 'B_PERSON', 'group_id'], as_index=False).agg({
        'DATE_TIME': 'min',
        'DURATION': 'sum'
    })

    # 时间格式还原为输入一致的带单引号字符串格式
    res_df['DATE_TIME'] = res_df['DATE_TIME'].dt.strftime("'%Y-%m-%d %H:%M:%S'")

    # 返回指定列，匹配输出格式要求
    return res_df[['A_PERSON', 'B_PERSON', 'DATE_TIME', 'DURATION']]

# 样例1测试
data1 = [
    [190,390,'\'2020-04-20 12:44:36\'',323],
    [282,811,'\'2020-04-06 11:12:24\'',25],
    [495,414,'\'2020-04-20 11:22:13\'',11],
    [827,158,'\'2020-04-30 13:27:22\'',22],
    [827,158,'\'2020-04-30 13:27:44\'',15]
]
df1 = pd.DataFrame(data1, columns=['A_PERSON','B_PERSON','DATE_TIME','DURATION'])
print(merger(df1))

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