CODEWITHSUNDEEP
pythonpandaspandas-groupby
AmirX
AmirX

Reputation: 2717

deleting groups based on a condition from a dataframe - pandas groupby

This is my dataframe:

df = pd.DataFrame({'sym': list('aaaaaabb'), 'order': [0, 0, 1, 1, 0, 1, 0, 1], 'key': [2, 2, 2, 2, 3, 3, 4, 4],
                   'vol': [1000, 1000, 500, 500, 100, 100, 200, 200]})

I add another column to it:

df['vol_cumsum'] = df.groupby(['sym', 'key', 'order']).vol.cumsum()

let`s define the problem like this (instead of words). Check this:

df.groupby(['sym', 'key', 'order']).vol_cumsum.last()

Now I want to omit the groups that their vol_cumsum ,according to the above groupby, do not match. In this case I want to omit the first group from my df. My desired df looks like this:

4    3      0   a   100         100
5    3      1   a   100         100
6    4      0   b   200         200
7    4      1   b   200         200

Upvotes: 1

Views: 38

Answers (1)

jezrael
jezrael

Reputation: 862581

Use GroupBy.transform with GroupBy.last for Series with same size like original DaatFrame, then create nw column by DataFrame.assign with GroupBy.all:

df['vol_cumsum'] = df.groupby(['sym', 'key', 'order']).vol.cumsum()
s = df.groupby(['sym', 'key', 'order']).vol_cumsum.transform('last')
mask = df.assign(new=df['vol_cumsum'].eq(s)).groupby(['sym', 'key', 'order'])['new'].transform('all')

df = df[mask]
print (df)
  sym  order  key  vol  vol_cumsum
4   a      0    3  100         100
5   a      1    3  100         100
6   b      0    4  200         200
7   b      1    4  200         200

Upvotes: 1

Related Questions