Find groups containing several rows matching condition

user_1356000 0
Viewed 54

Given the following dataframe

df = pd.DataFrame({'A':list('AAAAAABBBBBCCCCCC'),
                   'B':list('EFGHIJEGHJKGHJKEI')})

|    | A   | B   |
|---:|:----|:----|
|  0 | A   | E   |
|  1 | A   | F   |
|  2 | A   | G   |
|  3 | A   | H   |
|  4 | A   | I   |
|  5 | A   | J   |
|  6 | B   | E   |
|  7 | B   | G   |
|  8 | B   | H   |
|  9 | B   | J   |
| 10 | B   | K   |
| 11 | C   | G   |
| 12 | C   | H   |
| 13 | C   | J   |
| 14 | C   | K   |
| 15 | C   | E   |
| 16 | C   | I   |

I would like to find all the elements in A where B contains "G", and "H", and "I"

therefore, the result should be

|    | A   | B   |
|---:|:----|:----|
|  2 | A   | G   |
|  3 | A   | H   |
|  4 | A   | I   |
| 11 | C   | G   |
| 12 | C   | H   |
| 16 | C   | I   |

For the moment, I have found the following solution, but that seems overly hackish and I feel I'm missing something obvious

hit = list('GHI')
out = df[df.groupby('A').apply(lambda x: (x['B'].isin(hit))&(x['B'].isin(hit).sum()==len(hit))).values]
pandas python
5 Answers

This transformation isn't that obvious, but we can keep it vectorized by checking the size of each group compared to hit:

d = df[df['B'].isin(hit)]
size = d.groupby('A').size()
grps = size[size.eq(len(hit))].index

d[d['A'].isin(grps)]

    A  B
2   A  G
3   A  H
4   A  I
11  C  G
12  C  H
16  C  I
user_90812670

Let us do groupby + filter

hit = list('GHI')
out = df.groupby('A').filter(lambda x : pd.Series(hit).isin(x['B']).all())
out = out[out.B.isin(hit)]
out
Out[308]: 
    A  B
2   A  G
3   A  H
4   A  I
11  C  G
12  C  H
16  C  I
user_79645270

You have two conditions:

  1. hit is a subset of the group: x['B'].isin(hit).sum()==len(hit)

  2. the value at B is contained in hit: x['B'].isin(hit)

So you can express both conditions like this

hit = frozenset('GHI')
print(df[df.groupby('A')['B'].transform(hit.issubset) & df['B'].isin(hit)])

Output

    A  B
2   A  G
3   A  H
4   A  I
11  C  G
12  C  H
16  C  I

The expression:

df.groupby('A')['B'].transform(hit.issubset)

is the equivalent of condition 1.

user_40015920

Another way: First remove those not in hit, then filter those that have all hit

(df[df['B'].isin(hit)]
   .drop_duplicates(['A','B'])
   .loc[lambda x: x.groupby('A')['A'].transform('size')==len(hit)]
)

Or similar idea with groupby().filter:

(df[df['B'].isin(hit)]
   .groupby('A')
   .filter(lambda x: x['B'].nunique()==len(hit))
)

Output:

    A  B
2   A  G
3   A  H
4   A  I
11  C  G
12  C  H
16  C  I
user_42384080

Here is another method of doing this with boolean indexing without using groupby and using unstack -

df = pd.DataFrame({'A':list('AAAAAABBBBBCCCCCC'),
                   'B':list('EFGHIJEGHJKGHJKEI')})


filt = df['B'].isin(hit)

df['C']=1
df[filt & df['A'].isin(df[filt].set_index(['A','B']).unstack('B').dropna().index)]

    A   B   C
2   A   G   1
3   A   H   1
4   A   I   1
11  C   G   1
12  C   H   1
16  C   I   1
user_47559540
Related