Remove groups from a DataFrame that contain only a single unique value in one column

Viewed 148

I am processing data with Pandas. 'A' is a unique ID column and column 'E' contains either 1 or 0. I want to keep only groups where the value of column E contains both 0 and 1. (I want to delete rows where columns A are 2 and 4 as those groups contain only 1 and 0s respectively, leaving only rows where columns A are 1, 3, 5).

What is the best way to do this?

    A   B   C   D   E   F
1   1   0   0   0   1   1163.7
2   1   0.8 0.8 2.2 0   0
3   1   0.2 0.2 4.4 0   0
4   1   0.8 0.4 0.4 0   0
5   1   0.5 0.7 3.8 0   0
6   2   1   1   8.9 1   116
7   2   1.5 1.5 1.7 1   116
8   2   2   2   8.7 1   116
9   3   3   3   5.  0   0
10  3   4.5 4.5 2.2 0   0
11  3   6.0 6.5 0.8 0   0
12  3   8   8   0.3 0   0
13  3   5.3 0   0   1   116
14  3   0   0   0   1   116
15  4   0.8 0.8 1.1 0   0
16  4   0.2 0.5 3.4 0   0
17  4   0.4 0.8 3.2 0   0
18  4   0.7 0.5 3.0 0   0
19  5   1   1   1.5 0   0
20  5   1.5 1.5 1.7 0   0
21  5   2   2   7.9 1   116

I want to get the following data.

       A   B   C   D   E   F

1   1   0   0   0   1   1163.7
2   1   0.8 0.8 2.2 0   0
3   1   0.2 0.2 4.4 0   0
4   1   0.8 0.4 0.4 0   0
5   1   0.5 0.7 3.8 0   0
6   3   3   3   2.2 0   0
7   3   4.5 4.5 2.2 0   0
8   3   6.0 6.5 0.8 0   0
9   3   8   8   0.3 0   0
10  3   5.3 0   0   1   116
11  3   0   0   0   1   116
12  5   1   1   1.5 0   0
13  5   1.5 1.5 1.7 0   0
14  5   2   2   7.9 1   116

2 Answers

Use Series.groupby on column E and transform using any to create a boolean mask:

m = (df['E'].eq(0).groupby(df['A']).transform('any') &
     df['E'].eq(1).groupby(df['A']).transform('any'))
df1 = df[m]

Or another idea if column E consists only of zeros and ones,

m = df.groupby('A')['E'].nunique().eq(2)
df1 = df[df['A'].isin(m[m].index)]

Result:

print(df1)
    A    B    C    D  E       F
1   1  0.0  0.0  0.0  1  1163.7
2   1  0.8  0.8  2.2  0     0.0
3   1  0.2  0.2  4.4  0     0.0
4   1  0.8  0.4  0.4  0     0.0
5   1  0.5  0.7  3.8  0     0.0
9   3  3.0  3.0  5.0  0     0.0
10  3  4.5  4.5  2.2  0     0.0
11  3  6.0  6.5  0.8  0     0.0
12  3  8.0  8.0  0.3  0     0.0
13  3  5.3  0.0  0.0  1   116.0
14  3  0.0  0.0  0.0  1   116.0
19  5  1.0  1.0  1.5  0     0.0
20  5  1.5  1.5  1.7  0     0.0
21  5  2.0  2.0  7.9  1   116.0

you can use drop_duplicates on columns A and E and groupby.size to see where the group by A has 2 different elements as E is only 0 or 1. Then use the index where the size is equal to 2 like:

s = df[['A','E']].drop_duplicates().groupby('A').size()
df_ = df[df['A'].isin(s[s.eq(2)].index)].copy()
print(df_)
    A    B    C    D  E       F
1   1  0.0  0.0  0.0  1  1163.7
2   1  0.8  0.8  2.2  0     0.0
3   1  0.2  0.2  4.4  0     0.0
4   1  0.8  0.4  0.4  0     0.0
5   1  0.5  0.7  3.8  0     0.0
9   3  3.0  3.0  5.0  0     0.0
10  3  4.5  4.5  2.2  0     0.0
11  3  6.0  6.5  0.8  0     0.0
12  3  8.0  8.0  0.3  0     0.0
13  3  5.3  0.0  0.0  1   116.0
14  3  0.0  0.0  0.0  1   116.0
19  5  1.0  1.0  1.5  0     0.0
20  5  1.5  1.5  1.7  0     0.0
21  5  2.0  2.0  7.9  1   116.0
Related