Python Code For Extracting Duplicates Conditionally

Question

Let's say we have the following dataframe:

   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   

What would be the Python code for only extracting the rows such that at least one 'id' exists in both 'exp' and 'control'?

The following is the desired output:

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

The following code was used for creating this dataframe:

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'])

I tried the below code, but it only extracts all the duplicates in 'id' unconditionally. My guess is to add an if statement?

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

Answer 1

Try this. Base on your sample data column group having only 2 values exp and 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

For efficiency, instead of using lambda and filter, use transform and boolean indexing:

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

Output:

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

Timings:

%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)