如果用户的任何行包含一定的值，则选择所有用户的行

Question

我有病人的数据，日期，药物和诊断。每个病人都有一个独特的id ('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，而且意味着所有治疗：

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

我目前的解决方案：

1)获取包括药物“B”在内的行的所有pid

2)获取步骤1中包含pid的所有行。

这个解决方案的问题是，我需要用所有pid(百万)来做一个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 on 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的回答，但如果你说

我目前的解决方案是: 1)获得包含药物'B‘2的行的所有pid )从步骤1获得所有包含pid的行。

可以解决的问题比硬编码的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)]