如何在Pandas DataFrame中实现向量化操作并获取满足条件的患者集合？

原始代码与需求 #

原始数据代码 #

import pandas as pd

columns = ['S1', 'S2', 'S3', 'S4', 'S5']

df = pd.DataFrame({'Patient':['p1', 'p2', 'p3', 'p4', 'p5', 'p6', 'p7', 'p8', 'p8', 'p10'],
                   'S1':[0.7, 0.3, 0.5, 0.8, 0.9, 0.1, 0.9, 0.2, 0.6, 0.3],
                   'S2':[0.2, 0.3, 0.5, 0.4, 0.9, 0.1, 0.9, 0.7, 0.4, 0.3],
                   'S3':[0.6, 0.3, 0.5, 0.8, 0.9, 0.8, 0.9, 0.3, 0.6, 0.3],
                   'S4':[0.2, 0.3, 0.7, 0.8, 0.9, 0.1, 0.9, 0.7, 0.3, 0.3 ],
                   'S5':[0.9, 0.8, 0.5, 0.8, 0.9, 0.7, 0.2, 0.7, 0.6, 0.3 ]})

# 现有向量化操作示例
# 统计每一列中值≥0.5的单元格数量
arr1 = df[columns].ge(0.5).sum().to_numpy()

# 计算每一列中值≥0.5的单元格之和
arr2 = df[df[columns]>=0.5][columns].sum().to_numpy()

print(arr1)
print(arr2)

需求说明 #

为上述DataFrame中的S1至S5每一列，提取列值≥0.5对应的Patient列患者列表（或去重后的集合），输出格式参考如下：

[('p1', 'p3', 'p4', 'p5', 'p7', 'p8'), 
 ('p3', 'p5', 'p7', 'p8'), 
 ('p1', 'p3', 'p4', 'p5', 'p6', 'p7', 'p8'), 
 (...),
 (...)]

解决方案代码 #

import pandas as pd

columns = ['S1', 'S2', 'S3', 'S4', 'S5']

df = pd.DataFrame({'Patient':['p1', 'p2', 'p3', 'p4', 'p5', 'p6', 'p7', 'p8', 'p8', 'p10'],
                   'S1':[0.7, 0.3, 0.5, 0.8, 0.9, 0.1, 0.9, 0.2, 0.6, 0.3],
                   'S2':[0.2, 0.3, 0.5, 0.4, 0.9, 0.1, 0.9, 0.7, 0.4, 0.3],
                   'S3':[0.6, 0.3, 0.5, 0.8, 0.9, 0.8, 0.9, 0.3, 0.6, 0.3],
                   'S4':[0.2, 0.3, 0.7, 0.8, 0.9, 0.1, 0.9, 0.7, 0.3, 0.3 ],
                   'S5':[0.9, 0.8, 0.5, 0.8, 0.9, 0.7, 0.2, 0.7, 0.6, 0.3 ]})

# 按列提取满足条件的患者（去重后转元组）
patient_lists = [tuple(df.loc[df[col] >= 0.5, 'Patient'].unique()) for col in columns]

print(patient_lists)

输出结果 #

[('p1', 'p3', 'p4', 'p5', 'p7', 'p8'),
 ('p3', 'p5', 'p7', 'p8'),
 ('p1', 'p3', 'p4', 'p5', 'p6', 'p7', 'p8'),
 ('p3', 'p4', 'p5', 'p7', 'p8'),
 ('p1', 'p2', 'p3', 'p4', 'p5', 'p6', 'p8')]

代码说明 #

• 用列表推导式遍历每一列，通过df.loc[df[col] >= 0.5, 'Patient']筛选出当前列值≥0.5的患者；
• unique()用于去除重复患者（如原数据中p8出现两次）；
• 最后转为元组，与需求的输出格式一致。

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