Ari Frid
Ari Frid

Reputation: 127

Filtering a Pandas Dataframe by an aggregate function?

So I have dataframe that looks like this:

    STORE    PRODUCT    INVENTORY
1   store1      a           1
2   store1      b           0
3   store2      a           0
4   store2      b           0
5   store3      a           1
6   store3      b           1

I want to filter this such that it only shows me stores with a total inventory > 0 in the sense that the result will be as follows because store2 has a total inventory of 0.

    STORE    PRODUCT    INVENTORY
1   store1      a           1
2   store1      b           0
5   store3      a           1
6   store3      b           1

How do I do this?

Upvotes: 2

Views: 136

Answers (2)

ThomasIsCoding
ThomasIsCoding

Reputation: 102309

We can use filter like below

df.groupby('STORE').filter(lambda x: x['INVENTORY'].sum()>0)

which gives

    STORE   PRODUCT INVENTORY
1   store1  a   1
2   store1  b   0
5   store3  a   1
6   store3  b   1

Data

import pandas as pd
import io

text = """
    STORE    PRODUCT    INVENTORY
1   store1      a           1
2   store1      b           0
3   store2      a           0
4   store2      b           0
5   store3      a           1
6   store3      b           1
"""
df = pd.read_csv(io.StringIO(text), sep = '\s+')

Upvotes: 2

Anoushiravan R
Anoushiravan R

Reputation: 21938

You can try:

df.loc[(df.groupby('STORE')['INVENTORY'].transform(sum) > 0)]

    STORE PRODUCT  INVENTORY
1  store1       a          1
2  store1       b          0
5  store3       a          1
6  store3       b          1

Upvotes: 3

Related Questions