分组数据日期列对比：用iterrows实现患者出院日期区间计数

问题背景 #

现有如下DataFrame：

patient      admit       discharge     30d_admit
   0       a       2022-03-01   2022-03-01   2022-01-31
   1       a       2022-02-02   2022-02-02   2022-01-04
   2       a       2022-01-28   2022-01-28   2021-12-30
   3       b       2022-12-04   2022-12-04   2022-11-05
   4       c       2022-11-13   2022-11-13   2022-10-15
   5       c       2022-10-23   2022-10-23   2022-09-24

需求：按patient分组后，对每条admit记录，统计该患者所有discharge日期中落在当前记录的30d_admit与admit之间的数量，生成count列，预期输出：

patient      admit     discharge  30d_admit     count
   0        a          2022-03-01 2022-03-01 2022-01-31     1
   1        a          2022-02-02 2022-02-02 2022-01-04     1
   2        a          2022-01-28 2022-01-28 2021-12-30     0
   3        b          2022-12-04 2022-12-04 2022-11-05     0
   4        c          2022-11-13 2022-11-13 2022-10-15     1
   5        c          2022-10-23 2022-10-23 2022-09-24     0

解决方案 #

方法1：分组后使用iterrows()遍历 #

首先需将日期列转换为datetime类型，才能进行区间比较：

import pandas as pd

# 转换日期列为datetime类型
df['admit'] = pd.to_datetime(df['admit'])
df['discharge'] = pd.to_datetime(df['discharge'])
df['30d_admit'] = pd.to_datetime(df['30d_admit'])

按patient分组，遍历每组内的每条记录，统计符合条件的出院日期数量：

# 初始化count列
df['count'] = 0

# 分组处理每个患者的数据
for patient, group in df.groupby('patient'):
    # 提取当前患者所有出院日期
    all_discharges = group['discharge'].values
    # 遍历组内每条记录
    for idx, row in group.iterrows():
        # 筛选出落在[30d_admit, admit]区间的出院日期
        mask = (all_discharges > row['30d_admit']) & (all_discharges < row['admit'])
        # 统计数量并赋值
        df.loc[idx, 'count'] = mask.sum()

方法2：高效向量化实现（推荐） #

iterrows()遍历效率较低，针对大数据量可使用广播方式优化：

import pandas as pd

# 转换日期类型
df['admit'] = pd.to_datetime(df['admit'])
df['discharge'] = pd.to_datetime(df['discharge'])
df['30d_admit'] = pd.to_datetime(df['30d_admit'])

def calculate_count(group):
    # 广播生成比较矩阵，批量判断每个出院日期是否符合每条记录的区间
    discharges = group['discharge'].values[:, None]
    lower_bounds = group['30d_admit'].values
    upper_bounds = group['admit'].values
    # 按列统计符合条件的数量
    counts = ((discharges > lower_bounds) & (discharges < upper_bounds)).sum(axis=0)
    group['count'] = counts
    return group

# 分组应用计算逻辑
df = df.groupby('patient').apply(calculate_count).reset_index(drop=True)

结果说明 #

两种方法均可得到预期输出，比如患者a的第一条记录，其所有出院日期中仅2022-02-02落在2022-01-31与2022-03-01之间，因此count值为1。

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