如何按多规则条件筛选pandas DataFrame中的行

原始数据 #

import pandas as pd
from pandas import Timestamp

# 原始DataFrame
df = pd.DataFrame({
    "VisitDate": {
        0: Timestamp("2018-04-16 00:00:00"),
        1: Timestamp("2018-05-15 00:00:00"),
        2: Timestamp("2018-05-15 00:00:00"),
        3: Timestamp("2018-05-24 00:00:00"),
        4: Timestamp("2019-11-02 00:00:00"),
        5: Timestamp("2019-11-05 00:00:00"),
        6: Timestamp("2018-05-15 00:00:00"),
        7: Timestamp("2018-08-25 00:00:00"),
        8: Timestamp("2018-09-10 00:00:00"),
        9: Timestamp("2018-09-14 00:00:00"),
        10: Timestamp("2018-12-07 00:00:00"),
        11: Timestamp("2019-05-01 00:00:00"),
        12: Timestamp("2020-08-06 00:00:00"),
        13: Timestamp("2018-08-20 00:00:00"),
    },
    "FirstDate": {
        0: pd.NaT,
        1: Timestamp("2018-05-15 00:00:00"),
        2: pd.NaT,
        3: pd.NaT,
        4: pd.NaT,
        5: Timestamp("2019-11-05 00:00:00"),
        6: pd.NaT,
        7: pd.NaT,
        8: Timestamp("2018-09-10 00:00:00"),
        9: Timestamp("2018-09-10 00:00:00"),
        10: Timestamp("2018-09-10 00:00:00"),
        11: Timestamp("2018-09-10 00:00:00"),
        12: pd.NaT,
        13: pd.NaT,
    },
    "ID": {
        0: 7410013407,
        1: 7410013407,
        2: 7410013407,
        3: 7410013407,
        4: 7410022611,
        5: 7410022611,
        6: 8723458,
        7: 8723458,
        8: 8723458,
        9: 8723458,
        10: 8723458,
        11: 8723458,
        12: 8723458,
        13: 8723458,
    },
    "Dx": {
        0: "OtherPA",
        1: "PA",
        2: "OtherPA",
        3: "OtherPA",
        4: "OtherPA",
        5: "OM",
        6: "OtherPA",
        7: "OtherPA",
        8: "PA",
        9: "PA",
        10: "OM",
        11: "OM",
        12: "OtherOM",
        13: "OtherOM",
    },
})

# 数据类型
print(df.dtypes)
# VisitDate    datetime64[ns]
# FirstDate    datetime64[ns]
# ID                    int64
# Dx                   object
# dtype: object

筛选规则 #

• 所有Dx为PA或OM的行必须保留；
• 若ID对应的FirstDate行Dx为PA，则保留VisitDate在FirstDate前30天内且Dx为OtherPA的行；
• 若ID对应的FirstDate行Dx为OM，则保留VisitDate在FirstDate前30天内且Dx为OtherOM的行。

预期结果 #

# 筛选后的数据
print(final_result.reset_index(drop=True))
#    VisitDate FirstDate          ID       Dx
# 0 2018-04-16       NaT  7410013407  OtherPA
# 1 2018-05-15 2018-05-15  7410013407       PA
# 2 2018-05-15       NaT  7410013407  OtherPA
# 3 2019-11-05 2019-11-05  7410022611       OM
# 4 2018-08-25       NaT     8723458  OtherPA
# 5 2018-09-10 2018-09-10     8723458       PA
# 6 2018-09-14 2018-09-10     8723458       PA
# 7 2018-12-07 2018-09-10     8723458       OM
# 8 2019-05-01 2018-09-10     8723458       OM
# 9 2018-08-20       NaT     8723458  OtherOM

实现代码 #

步骤1：提取每个ID的基准信息 #

首先获取每个ID对应的FirstDate，以及该FirstDate所在行的Dx（作为基准诊断）：

# 提取每个ID的FirstDate和对应的基准Dx
id_base_info = df.dropna(subset=['FirstDate']).groupby('ID').agg(
    base_first_date=('FirstDate', 'first'),
    base_dx=('Dx', 'first')
).reset_index()

# 将基准信息合并到原DataFrame，让每一行都能获取到所属ID的基准数据
df_with_base = df.merge(id_base_info, on='ID', how='left')

步骤2：构建筛选条件并过滤数据 #

根据规则构建三个筛选条件，用逻辑或连接后过滤：

# 条件1：保留Dx为PA或OM的行
condition1 = df_with_base['Dx'].isin(['PA', 'OM'])

# 条件2：基准Dx为PA，当前行是OtherPA，且VisitDate在基准FirstDate前30天内
condition2 = (df_with_base['base_dx'] == 'PA') & \
             (df_with_base['Dx'] == 'OtherPA') & \
             (df_with_base['VisitDate'] >= df_with_base['base_first_date'] - pd.Timedelta(days=30))

# 条件3：基准Dx为OM，当前行是OtherOM，且VisitDate在基准FirstDate前30天内
condition3 = (df_with_base['base_dx'] == 'OM') & \
             (df_with_base['Dx'] == 'OtherOM') & \
             (df_with_base['VisitDate'] >= df_with_base['base_first_date'] - pd.Timedelta(days=30))

# 合并条件，筛选数据并移除临时列
final_result = df_with_base[condition1 | condition2 | condition3].drop(columns=['base_first_date', 'base_dx'])

执行上述代码后，final_result即为符合要求的数据集，与预期结果一致。

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