Python3中按条件筛选Pandas DataFrame行（需含MDR与ICD）

原始DataFrame #

CUI         id          term id term_name 
C0000729    10000057    MDR LLT Abdominal cramps
C0000729    10000056    MDR LLT Abdominal cramp
C0000729    10011286    MDR LLT Cramp abdominal
C0000729    10000058    MDR LLT Abdominal crampy pains
C0000729    10093764    ICD10   PT  Abdominal crampy pains
C0000800    10093765    ICD10   PT  Abdominal pain
C0000800    10000058    MDR LLT Abdominal crampy pains
C0000800    10093764    ICD10AM PT  Abdominal crampy pains
C0000730    10000052    MDR LLT Abdominal cramps back

需求说明 #

仅保留每个CUI下同时包含MDR和至少一种**ICD类（ICD10、ICD10AM）**的所有记录，排除仅含MDR的CUI（如最后一行的C0000730）。

预期输出 #

CUI         id          term id term_name 
C0000729    10000057    MDR LLT Abdominal cramps
C0000729    10000056    MDR LLT Abdominal cramp
C0000729    10011286    MDR LLT Cramp abdominal
C0000729    10000058    MDR LLT Abdominal crampy pains
C0000729    10093764    ICD10   PT  Abdominal crampy pains
C0000800    10093765    ICD10   PT  Abdominal pain
C0000800    10000058    MDR LLT Abdominal crampy pains
C0000800    10093764    ICD10AM PT  Abdominal crampy pains

用户原有代码（存在问题） #

#select only those mappings where ICD and MeDRA both exists for a particular CUI id
s = set(['ICD10','ICD10CM','ICD10AM','MDR'])

dff_mapped = df_umls[df_umls.groupby('CUI')['SAB'].transform(lambda x: set(x) == s)]

dff_mapped = dff_mapped.sort_values(['CUI', 'SAB'],ascending = [True, True])

dff_mapped.to_csv('df_mapped', index = False, sep = ',')

代码修正及解释 #

原有代码的问题在于：要求每个CUI必须包含集合s中的所有4种类型（ICD10、ICD10CM、ICD10AM、MDR），但实际需求只需要同时存在MDR和至少一种ICD类即可，不需要覆盖所有ICD子类。

修正后的代码如下：

# 定义ICD类的集合和必需的MDR类型
icd_types = {'ICD10', 'ICD10CM', 'ICD10AM'}
required_mdr = {'MDR'}

# 对每个CUI判断：同时包含MDR，且至少有一个ICD类记录
filter_mask = df_umls.groupby('CUI')['SAB'].transform(
    lambda x: required_mdr.issubset(x) and not icd_types.isdisjoint(x)
)

# 筛选符合条件的行并排序
dff_mapped = df_umls[filter_mask].sort_values(['CUI', 'SAB'], ascending=[True, True])

# 保存结果（注意添加.csv后缀）
dff_mapped.to_csv('df_mapped.csv', index=False, sep=',')

关键逻辑说明 #

• required_mdr.issubset(x)：确保该CUI下存在MDR类型的记录
• not icd_types.isdisjoint(x)：确保该CUI下至少存在一种ICD类的记录（与ICD集合有交集）
• 两者同时满足时，该CUI的所有行都会被保留，符合需求。

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