有条件地提取重复项的Python代码

Question

假设我们有以下数据帧：

   group    id  performance
0   exp      A   1
1   exp      B   2
2   exp      B   3
3   exp      C   4
4   control  A   5
5   control  A   6
6   exp      D   7
7   control  D   8   

如果Python代码只提取行，使得exp和control中至少有一个id存在，那该怎么办？你知道吗

以下是所需的输出：

   group    id  performance
0   exp      A   1
4   control  A   5
5   control  A   6
6   exp      D   7
7   control  D   8  

以下代码用于创建此数据帧：

students = [('exp', 'A', 1),
            ('exp', 'B', 2),
            ('exp', 'B', 3),
            ('exp', 'C', 4),
            ('control', 'A', 5),
            ('control', 'A', 6),
            ('exp', 'D', 7),
            ('control', 'D', 8)
            ]

import pandas as pd
student_df = pd.DataFrame(students, columns=['group', 'id', 'performance'])

我尝试了下面的代码，但它只能无条件地提取“id”中的所有重复项。我猜应该加一个if语句？你知道吗

student_df[student_df.duplicated(['id'], keep = False)]

Answer 1

试试这个。基于示例数据列group只有2个值exp和control：

df_out = student_df.groupby('id').filter(lambda x: x.group.nunique() > 1)

Out[570]:
     group id  performance
0      exp  A            1
4  control  A            5
5  control  A            6
6      exp  D            7
7  control  D            8

Answer 2

为了提高效率，请使用transform和布尔索引，而不是使用lambda和filter：

student_df[student_df.groupby('id')['group'].transform('nunique')>1]

输出：

     group id  performance
0      exp  A            1
4  control  A            5
5  control  A            6
6      exp  D            7
7  control  D            8

时间安排：

%timeit student_df.groupby('id').filter(lambda x: x.group.nunique() > 1)
5.29 ms ± 165 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

%timeit student_df[student_df.groupby('id')['group'].transform('nunique')>1]
3.01 ms ± 113 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)