I have a pandas dataframe that looks like this:
import pandas as pd
d = {'group': ['A', 'A', 'A', 'A', 'A', 'B', 'B', 'B', 'B', 'B', 'C', 'C', 'C', 'C', 'C'],
'fruit': ['strawberry', 'strawberry', 'strawberry', 'strawberry', 'strawberry', 'kiwi', 'banana', 'strawberry', 'orange', 'kiwi', 'banana', 'melon', 'orange', 'kiwi', 'melon']}
df = pd.DataFrame(data=d)
df
group fruit
A strawberry
A strawberry
A strawberry
A strawberry
A strawberry
B kiwi
B banana
B strawberry
B orange
B kiwi
C banana
C melon
C orange
C kiwi
C melon
Using the code below i'm able to remove a group
if that group contains any row with the fruit
value 'strawberry'.
df[~df.fruit.eq('strawberry').groupby(df.group).transform('any')]
group fruit
C banana
C melon
C orange
C kiwi
C melon
However, I want to only remove a group
if that group contains at least two rows with the fruit value
'strawberry', meaning that the end result should als include group B. How do I achieve this?
CodePudding user response:
Change to transform('sum')
n = 2
out = df[~(df.fruit.eq('strawberry').groupby(df.group).transform('sum')>=n)]
out
Out[108]:
group fruit
5 B kiwi
6 B banana
7 B strawberry
8 B orange
9 B kiwi
10 C banana
11 C melon
12 C orange
13 C kiwi
14 C melon