Home > database >  Remove group from the pandas dataframe when a specific value within the group occurs at least two ti
Remove group from the pandas dataframe when a specific value within the group occurs at least two ti

Time:11-17

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
  • Related