如何筛选pid_code与mo一致且采样间隔≤2天的DataFrame行

我有一个包含4列的DataFrame：

• 微生物数据(mo)
• 采样日期(sample_date)
• 患者ID(pid_code)
• 每行唯一标识(id)

数据示例 #

需求 #

筛选出pid_code与mo均相同，且组内存在至少两个采样日期间隔不超过2天的行，期望结果如下：

期望结果 #

步骤1：统一日期格式 #

首先将sample_date转换为datetime类型，确保能正确计算日期差：

import pandas as pd

# 构建示例DataFrame（实际使用时替换为你的数据读取逻辑）
df = pd.DataFrame({
    'id': [1,2,3,4,5,6,7],
    'mo': ['Staphylococcus epidermidis', 'Staphylococcus epidermidis', 'Staphylococcus capitis', 'Staphylococcus hominis', 'Staphylococcus hominis', 'Micrococcus luteus', 'Micrococcus luteus'],
    'sample_date': ['2019-12-13', '2019-12-14', '2020-01-13', '2021-11-14', '2021-11-17', '2023-06-06', '2023-06-07'],
    'pid_code': [120,120,131,155,155,133,133]
})

# 转换为datetime类型
df['sample_date'] = pd.to_datetime(df['sample_date'])

步骤2：标记符合条件的分组 #

按pid_code和mo分组，检查每组内是否存在间隔≤2天的采样日期：

def check_interval(group):
    # 按日期排序后计算相邻日期差
    sorted_group = group.sort_values('sample_date')
    date_diff = sorted_group['sample_date'].diff().dt.days
    # 判断是否有满足条件的日期差
    return any(date_diff.dropna() <= 2)

# 获取所有符合条件的(pid_code, mo)组合
valid_pairs = df.groupby(['pid_code', 'mo']).apply(check_interval)
valid_pairs = valid_pairs[valid_pairs].index

步骤3：筛选目标行 #

根据标记的有效分组，提取对应的行数据：

# 筛选结果
result = df[df.set_index(['pid_code', 'mo']).index.isin(valid_pairs)]
# 按id排序保持原数据顺序（可选）
result = result.sort_values('id').reset_index(drop=True)

print(result)

运行输出 #

id                          mo sample_date  pid_code
0   1  Staphylococcus epidermidis  2019-12-13       120
1   2  Staphylococcus epidermidis  2019-12-14       120
2   6          Micrococcus luteus  2023-06-06       133
3   7          Micrococcus luteus  2023-06-07       133

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