如何在Python Pandas中无循环实现基于事件与时间范围的分组？

输入数据 #

构造DataFrame代码 #

import pandas as pd

id_ = ['A','A','A','A','A','B','B','B','B']
days = [1,10,20,25,100,10,20,30,50]
event = [0,0,0,1,0,0,0,1,0]
df = pd.DataFrame({'id': id_, 'days': days, 'event': event})

原始数据预览 #

id  days  event
0  A     1      0
1  A    10      0
2  A    20      0
3  A    25      1
4  A   100      0
5  B    10      0
6  B    20      0
7  B    30      1
8  B    50      0

需求说明 #

按id分组，将event=1的行，以及该行日期前后30天内的所有同ID行归为同一组；不允许使用for循环，最终输出需新增group列，格式符合预期。

解决方案 #

实现思路 #

借助Pandas的分组与广播特性批量处理：

1. 按id拆分数据组，提取每组内触发事件的日期，计算每个事件的有效时间区间（日期±30天）
2. 对组内每一行的日期，判断是否落在任意一个事件的有效区间内
3. 给符合条件的行分配组号（按ID顺序递增），不符合的标记为NaN

代码实现 #

import pandas as pd

# 构造数据
id_ = ['A','A','A','A','A','B','B','B','B']
days = [1,10,20,25,100,10,20,30,50]
event = [0,0,0,1,0,0,0,1,0]
df = pd.DataFrame({'id': id_, 'days': days, 'event': event})

# 将id转为类别类型，方便生成连续组号
df['id'] = df['id'].astype('category')

def process_group(group):
    # 获取当前组内所有事件触发的日期
    event_days = group[group['event'] == 1]['days'].values
    if not len(event_days):
        group['group'] = pd.NA
        return group
    # 检查每行days是否在任意事件的±30天区间内
    in_range = ((group['days'] >= (event_days - 30)) & (group['days'] <= (event_days + 30))).any(axis=1)
    # 分配组号：同一ID组的有效行共用一个组号，按ID顺序从1开始递增
    group['group'] = in_range.astype(float) * (group['id'].cat.codes.iloc[0] + 1)
    return group

# 分组处理并合并结果
result = df.groupby('id', group_keys=False).apply(process_group)

print(result)

输出结果 #

id  days  event  group
0  A     1      0    1.0
1  A    10      0    1.0
2  A    20      0    1.0
3  A    25      1    1.0
4  A   100      0    NaN
5  B    10      0    2.0
6  B    20      0    2.0
7  B    30      1    2.0
8  B    50      0    2.0

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