Pandas中条件计数向量化优化方案咨询

嘿，我懂你想通过向量化提升30天再入院次数统计性能的痛点——df.rolling().apply确实经常因为不是真正的向量化执行，在数据量大的时候拖慢速度。咱们从数据预处理开始，一步步给出高效的解决方案。

第一步：数据预处理 #

首先得把你的模拟数据转换成规范的DataFrame，重点是把日期列转成datetime类型，并且按患者ID和就诊日期排序，确保后续计算的时间顺序正确：

import pandas as pd

# 模拟你的原始数据
data = {
    'ID': [1, 1, 2, 3, 3],
    'VISIT_NO': [1, 2, 1, 1, 2],
    'ARRIVED': ['29/02/1996', '01/12/1996', '20/09/1996', '27/06/1996', '04/07/1996'],
    'LEFT': ['01/03/1996', '04/12/1996', '21/09/1996', '28/06/1996', '06/07/1996'],
    'HAD_A_MASSAGE': [0, 1, 1, 1, 0],
    'BROUGHT_A_FRIEND': [1, 0, 0, 0, 1]
}

df = pd.DataFrame(data)

# 转换日期列（注意你的日期格式是DD/MM/YYYY）
df['ARRIVED'] = pd.to_datetime(df['ARRIVED'], format='%d/%m/%Y')
df['LEFT'] = pd.to_datetime(df['LEFT'], format='%d/%m/%Y')

# 按ID和就诊日期排序，保证每个患者的记录是时间顺序
df = df.sort_values(['ID', 'ARRIVED']).reset_index(drop=True)

第二步：向量化统计再入院次数 #

方案1：分组广播运算（推荐，性能最优） #

这个方法利用NumPy的广播机制，对每个患者的所有就诊记录做批量时间差计算，完全避免循环：

def count_readmissions(group):
    # 将出院日期转成列向量，入院日期转成行向量，做广播计算时间差
    left_dates = group['LEFT'].values.reshape(-1, 1)
    arrived_dates = group['ARRIVED'].values
    
    # 筛选：后续入院日期必须在当前出院日期之后，且间隔≤30天
    time_diff = arrived_dates - left_dates
    valid_readmissions = (time_diff > pd.Timedelta(0)) & (time_diff <= pd.Timedelta(days=30))
    
    # 统计每行符合条件的次数
    return valid_readmissions.sum(axis=1)

# 分组应用函数，得到每个就诊的30天再入院次数
df['30DAY_READMISSION_COUNT'] = df.groupby('ID', group_keys=False).apply(count_readmissions)

方案2：IntervalIndex 区间匹配 #

利用Pandas的IntervalIndex快速匹配落在30天窗口内的入院记录，逻辑更直观：

def count_readmissions_interval(group):
    # 为每个出院日期创建[LEFT, LEFT+30天)的时间区间
    intervals = pd.IntervalIndex.from_arrays(
        group['LEFT'], 
        group['LEFT'] + pd.Timedelta(days=30), 
        closed='left'  # 包含出院当天，排除30天后的日期
    )
    
    # 统计每个入院日期落在哪些区间里，然后按区间求和
    return intervals.get_indexer(group['ARRIVED']).reshape(-1, len(intervals)).sum(axis=0)

df['30DAY_READMISSION_COUNT'] = df.groupby('ID', group_keys=False).apply(count_readmissions_interval)

为什么rolling.apply效果不好？ #

df.rolling().apply本质上是逐窗口循环执行自定义函数，没有利用Pandas/NumPy的向量化优化，当数据量较大时，性能会急剧下降。而上面的两种方案都是基于底层的向量化运算，能把性能提升几倍甚至几十倍。

结果验证 #

运行上面的代码后，你的DataFrame会新增30DAY_READMISSION_COUNT列，对应每个就诊记录的30天内再入院次数：

• ID=3的第一次就诊（LEFT=1996-06-28），后续6天就有一次入院，所以计数为1
• ID=1的两次就诊间隔超过30天，所以第一次计数为0，第二次没有后续就诊，计数为0

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