基于Pandas按地址+时间窗分组聚合相近事件的实现方案

需求描述 #

需要将时间相近且指定变量（如地址）取值相同的事件分组聚合。典型场景比如结合疾病发病日期与地址，识别同一地点、指定时间范围内的疾病暴发事件，需处理包含30万行数据的Pandas DataFrame。

示例数据 #

import pandas as pd

df = pd.DataFrame(
    [
        ['2020-01-01 10:00', '1', 'A'],
        ['2020-01-01 10:01', '2', 'A'],
        ['2020-01-01 10:02', '3a', 'A'],
        ['2020-01-01 10:02', '3b', 'A'],
        ['2020-01-02 10:03', '4', 'B'],
        ['2020-01-02 10:50', '5', 'B'],
        ['2020-01-02 10:54', '6', 'B'],
        ['2020-01-02 10:55', '7', 'B'],
    ], columns=['event_time', 'event_id', 'Address']
)

输出要求 #

输出需包含每组事件的起始时间、结束时间、事件ID列表及对应地址，格式示例如下：

event_time_start  event_time_end     events_and_related_event_id_list  Address
0 2020-01-01 10:00:00  2020-01-01 10:02:00        [1, 2, 3a]     A
6 2020-01-01 10:54:00  2020-01-01 10:55:00        [6, 7]         B

优化后解决方案 #

初始方案基于时间窗分组思路，但实际数据存在非连续、非序列事件，因此做了针对性优化，完整代码如下：

import pandas as pd

df = pd.DataFrame(
    [
        ['1', 'A', '2020-01-01 10:00'],
        ['2', 'B', '2020-01-01 10:01'],
        ['3', 'A', '2020-01-01 10:01'],
        ['4', 'C', '2020-01-01 10:02'],
        ['5', 'D', '2020-01-01 10:03'],
        ['6', 'A', '2020-01-01 10:03'],
        ['7', 'E', '2020-01-01 10:03'],
        ['8', 'A', '2020-01-01 10:07'],
        ['9', 'A', '2020-01-01 10:09'],
        ['10', 'A', '2020-01-01 10:11'],
        ['11', 'F', '2020-01-01 10:54'],
        ['12', 'G', '2020-01-01 10:55'],
        ['13', 'F', '2020-01-01 10:56'],
    ], columns=['id', 'Address', 'event_time']
)

# 按地址和时间排序，确保同地址事件按时间顺序处理
df = df.sort_values(by=["Address", "event_time"])
df['event_time'] = pd.to_datetime(df['event_time'])

# 第一步：按地址分组，标记每个事件的时间相近关联ID
timeDiff = pd.Timedelta("2m")  # 定义时间相近的阈值（2分钟）
def idsNearDates(mDf):
    # 找出当前事件时间前后timeDiff范围内的所有其他事件ID
    f = lambda colName, val: mDf.loc[mDf['event_time'].between(val - timeDiff, val + timeDiff),
                            'id'].drop(colName).tolist()
    mDf['relatedIds'] = [f(colName, value) for colName, value in mDf['event_time'].items()]
    return mDf
df_1stStep = df.groupby('Address').apply(idsNearDates).sort_values(by=["Address", 'event_time'])

# 第二步：将关联事件聚合为单一行
# 标记时间间隔超过阈值的位置，用于分组
timeGroup_bool = ~(df_1stStep['event_time'].between(df_1stStep['event_time'].shift(1) - timeDiff,
                                                             df_1stStep['event_time'].shift(1) + timeDiff))
# 合并同一组内的所有关联ID并去重
f1 = lambda x: list(dict.fromkeys([value for idList in x for value in idList]))
df_2ndstep = (df_1stStep.groupby([(timeGroup_bool).cumsum(),'Address'])
                   .agg(Date_first=('event_time','min'),
                        Date_last=('event_time','max'),
                        Ids=('relatedIds',f1))
                   .droplevel(0)
                   .reset_index())
# 过滤掉无关联事件的行
df_2ndstep = df_2ndstep[df_2ndstep['Ids'].str.len() > 0]

方案说明 #

1. 先按地址和时间排序，保证同地址的事件按时间顺序处理，避免逻辑混乱。
2. 按地址分组后，为每个事件标记出时间阈值内的所有关联事件ID。
3. 通过时间间隔判断生成分组标识，将连续时间范围内的事件合并为一组，同时对关联ID去重。
4. 最后过滤掉没有关联事件的行，得到最终的聚合结果。

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