如何高效计算DataFrame中患者首次患病与新增无病诊断的时间差

问题背景 #

现有一份已排序的面板数据DataFrame，字段包含ID、Date、diagnosisCode、hasDisease，其中诊断编码与患病标签对应关系定义如下：

positiveDiagnosisCodes = ['A']  # 对应hasDisease = 1（患病）
negativeDiagnosisCodes = ['B', 'C', 'D']  # 对应hasDisease = 0（无病）

需要计算每个患者两类时间的差值：

1. 该患者首次hasDisease=1的时间
2. 该患者首次出现未记录过的negativeDiagnosisCodes编码对应的hasDisease=0的时间

示例结果：

• ID2的时间差为2天
• ID3对应编码C、D的时间差分别为18天、23天

高效实现方案 #

无需复杂数据结构，用Pandas原生操作即可完成，步骤如下：

1. 预处理日期字段 #

确保Date为datetime类型，否则无法计算时间差：

import pandas as pd

df['Date'] = pd.to_datetime(df['Date'])

2. 提取每个患者首次患病时间 #

利用分组聚合直接取最小值（数据已排序，最小值即首次出现时间）：

# 筛选患病记录，按ID分组取最早日期
first_positive = df[df['hasDisease'] == 1].groupby('ID')['Date'].min().rename('first_positive_date')

3. 提取每个患者首次出现的未记录阴性诊断时间 #

通过分组遍历，用集合跟踪已出现过的阴性编码，筛选出每个编码的首次出现记录：

def extract_first_negative(group):
    # 筛选当前患者的阴性诊断记录
    negative_records = group[group['diagnosisCode'].isin(negativeDiagnosisCodes)]
    seen_codes = set()
    first_negative_list = []
    
    # 遍历已排序的记录，取每个阴性编码的首次出现
    for _, row in negative_records.iterrows():
        code = row['diagnosisCode']
        if code not in seen_codes:
            seen_codes.add(code)
            first_negative_list.append({
                'ID': row['ID'],
                'diagnosisCode': code,
                'first_negative_date': row['Date']
            })
    return pd.DataFrame(first_negative_list)

# 分组处理后重置索引
first_negative_df = df.groupby('ID').apply(extract_first_negative).reset_index(drop=True)

4. 合并数据并计算时间差 #

将首次患病时间与首次阴性诊断记录合并，计算天数差值：

# 合并两个数据集
result = first_negative_df.merge(first_positive, on='ID', how='left')
# 计算时间差（单位：天）
result['days_diff'] = (result['first_negative_date'] - result['first_positive_date']).dt.days

示例输出 #

处理后结果格式如下：

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