从数据框中筛选行，并创建仅具有相似值的新数据框

Question

我正在尝试取一个"患者"one_answers"药物"的数据框架。数据，并创建一个新的数据框架，仅包含使用过两种药物的患者。

import pandas as pd
data = [['1001', 'A'], ['1001', 'B'], ['1002', 'A'], ['1003', 'B'], ['1004', 'A'], ['1004', 'B']]
df = pd.DataFrame(data, columns = ['patient', 'drug'])

patient drug
1001    A
1001    B
1002    A
1003    B
1004    A
1004    B

我首先对drug值进行过滤并创建2个独立的数据框架。

df1 =df.loc[(df['drug'] == 'A')]
df2 =df.loc[(df['drug'] == 'B')]

我的问题是，我想结合这2个数据框，但只保留行有相同的病人在两个数据框。我试着使用。concat，但它没有给我我想要的，我只是把数据框组合在一起。

pd.concat([df1, df2], join="inner")

这就是我想要的输出

patient drug
1001    A
1001    B
1004    A
1004    B

我知道我可能只是错过了一个简单的参数，但我不知道它是什么。

Answer 1

使用groupby:

>>> df.loc[df.groupby("patient")["drug"].transform("nunique")==2]
patient drug
0    1001    A
1    1001    B
4    1004    A
5    1004    B

Answer 2

还可以在patient列的2个dfs上与.merge()进行内归并，得到df1和df2中同时出现的患者列表。然后，使用.loc和isin()对原df中的患者进行筛选，如下:

df.loc[df['patient'].isin(df1.merge(df2, how='inner', on='patient')['patient'])]

结果:

patient drug
0    1001    A
1    1001    B
4    1004    A
5    1004    B

附加测试用例

如果，如果原始df有' a '和'B'以外的第三种药物，例如

输入>

data = {'patient': [1001, 1001, 1001, 1002, 1003, 1004, 1004, 1004],
'drug': ['A', 'B', 'C', 'A', 'B', 'A', 'B', 'C']}
df = pd.DataFrame(data)
df
patient drug
0     1001    A
1     1001    B
2     1001    C
3     1002    A
4     1003    B
5     1004    A
6     1004    B
7     1004    C

df1 =df.loc[(df['drug'] == 'A')]
df2 =df.loc[(df['drug'] == 'B')]
df.loc[df['patient'].isin(df1.merge(df2, how='inner', on='patient')['patient'])]
patient drug
0     1001    A
1     1001    B
2     1001    C
5     1004    A
6     1004    B
7     1004    C

这里，我们还显示第三个药物'C'的行，因为我们仍然满足

的要求

仅保留两个数据框中具有相同患者的行。

(虽然使用了额外的药物，但仍然是同一组患者)。

如果，在这种情况下，我们想要一个干净的输出来显示只有药物' a '和'B'的条目，我们可以使用:

df_A_B = df1.append(df2).sort_values(['patient', 'drug'])
df_A_B.loc[df_A_B['patient'].isin(df1.merge(df2, how='inner', on='patient')['patient'])]

结果:

patient drug
0     1001    A
1     1001    B
5     1004    A
6     1004    B