pandas DataFrame按人逐行检查未来15天is_hosp为True的慢代码优化咨询

原代码性能瓶颈说明 #

你写的双层循环实现，本质是对每条观测都做一次全组范围的全量扫描，时间复杂度达到O(n²)，在百万行规模下，光是重复筛选子表、逐行Python循环的开销就能让任务跑几小时。我们完全可以利用“只需要找最近一次未来住院记录”的等价逻辑，把复杂度降到O(n log n)，实现秒级跑完。

性能最优方案：merge_asof 无循环实现 #

这个方案底层为C语言实现，无Python层循环开销，千万行规模下也可稳定运行，且逻辑与需求完全等价：

1. 先对全表按person_id、dates做升序排序，这是所有匹配逻辑的基础

   df = df.sort_values(['person_id', 'dates']).reset_index(drop=True)
   

2. 单独提取所有is_hosp=True的记录构造辅助表

   # 重命名住院日期列，避免和原表日期冲突
   hosp_records = df.loc[df['is_hosp'], ['person_id', 'dates']].rename(columns={'dates': 'nearest_future_hosp_date'})
   

3. 用merge_asof做按组的向前匹配：对每一行观测，找到同person_id下、日期大于等于当前观测日期的最近一条住院记录

   df = pd.merge_asof(
       df,
       hosp_records,
       on='dates', # 以日期列作为匹配的排序键
       by='person_id', # 严格在同一个人组内匹配，不跨人找住院记录
       direction='forward' # 关键逻辑：向前查找（即找当前日期之后的记录）
   )
   

4. 计算日期间隔，生成最终的判断列

   # 计算最近的未来住院日和当前观测日的间隔，填充无匹配场景的空值为False
   df['is_inh_15d'] = (df['nearest_future_hosp_date'] - df['dates']).dt.days.le(15).fillna(False)
   # 删除临时的住院日期列
   df = df.drop(columns=['nearest_future_hosp_date'])
   

等价逻辑说明 #

这个方案和你原需求的判断逻辑完全对齐，没有任何偏差：

• 对每条观测行，只要未来15天内（含观测当天）存在任意一条住院记录，那么这条观测和最近一条未来住院记录的间隔必然≤15天
• 如果最近的未来住院记录和当前观测的间隔超过15天，那么更远的住院记录间隔只会更大，不可能落在15天范围内
• 严格按person_id分组匹配，不会把其他用户的住院记录错配到当前用户上
• 兼容每个用户观测日期不连续的场景，不需要强制日期为连续序列

次优方案：分组二分查找实现 #

如果因为版本限制没法使用merge_asof，可以用分组结合np.searchsorted的二分查找实现，性能同样远高于原双层循环方案：

def check_group_hosp(group):
    # 组内已按日期排序，直接提取所有住院日期的有序数组
    hosp_date_arr = group.loc[group['is_hosp'], 'dates'].to_numpy()
    # 无住院记录的组直接返回全False
    if len(hosp_date_arr) == 0:
        group['is_inh_15d'] = False
        return group
    # 用二分查找定位每个观测点对应的最近未来住院记录索引
    row_date_arr = group['dates'].to_numpy()
    hosp_idx = np.searchsorted(hosp_date_arr, row_date_arr)
    # 处理索引越界（即当前观测后无任何住院记录）的场景
    valid_mask = hosp_idx < len(hosp_date_arr)
    day_gap = np.full(len(group), np.nan)
    day_gap[valid_mask] = (hosp_date_arr[hosp_idx[valid_mask]] - row_date_arr[valid_mask]) / np.timedelta64(1, 'D')
    group['is_inh_15d'] = day_gap <= 15
    return group

df = df.groupby('person_id', group_keys=False).apply(check_group_hosp)

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