Pandas数据帧 - 如果他的任何行包含特定值，则选择所有用户的行

Question

我有一个包含患者，日期，药物和诊断的数据框。每位患者都有一个独特的身份（'pid'），可能会或可能不会用不同的药物治疗。

选择所有患者在某些时候接受某种药物治疗的最佳做法是什么？由于我的数据集非常庞大，for循环和if语句是最后的选择。

例：

IN:
pid drug
1   A
1   B
1   C
2   A
2   C
2   E
3   B
3   C
3   D
4   D
4   E
4   F

选择所有患者在某一时间接受药物'B'治疗。请注意，必须包括该患者的所有条目，这不仅意味着用药物B治疗，而且包括所有治疗：

OUT:
1   A
1   B
1   C
3   B
3   C
3   D

我目前的解决方案 1）获取包含药物'B'的行的所有pid 2）从步骤1获取包含pid的所有行。 这个解决方案的问题是我需要用所有pid（数百万）制作一个loooong if语句

Answer 1

这是一种方式。

s = df.groupby('drug')['pid'].apply(set)

result = df[df['pid'].isin(s['B'])]

#    pid drug
# 0    1    A
# 1    1    B
# 2    1    C
# 6    3    B
# 7    3    C
# 8    3    D

说明

• 创建映射系列s作为单独的初始步骤，以便它不需要重新计算每个结果。
• 对于比较，使用set进行O（1）复杂性查找。

Answer 2

最简单的方法涉及groupby + transform：

df[df.drug.eq('B').groupby(df.pid).transform('any')]

   pid drug
0    1    A
1    1    B
2    1    C
6    3    B
7    3    C
8    3    D

为了寻求更快的解决方案，请在groupby上调用df，而不是系列：

df[~df.groupby('pid').drug.transform(lambda x: x.eq('B').any())]

    pid drug
3     2    A
4     2    C
5     2    E
9     4    D
10    4    E
11    4    F

Answer 3

我支持COLDSPEED的答案，但如果你说

My current solution:
1) Get all pid for rows that includes drug 'B'
2) Get all rows that include pid from step 1.
Problem with this solution is that I need to make a loooong if-statement with all pid's (millions)

可以解决比硬编码if更简单的问题

patients_B = df.loc[df['drug'] == 'B', 'pid]

要么

patients_B = set(df.loc[df['drug'] == 'B', 'pid])

然后

result = df[df['pid'].isin(patients_B)]

Answer 4

IIUC filter

df.groupby('pid').filter(lambda x : (x['drug']=='B').any())
Out[18]: 
   pid drug
0    1    A
1    1    B
2    1    C
6    3    B
7    3    C
8    3    D