熊猫数据框-如果用户的任何行包含特定值，则选择所有用户的行

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（百万）创建一个很长的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作为一个单独的初始步骤，这样就无需为每个结果重新计算。
• 为了进行比较，请使用setO（1）复杂度查找。