удаление групп на основе условия из фрейма данных - pandas groupby

Question

Это мой фрейм данных:

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]})

Я добавляю к нему еще один столбец:

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

Давайте определим проблему следующим образом (вместо слов).Проверьте это:

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

Теперь я хочу опустить группы, которые не соответствуют их vol_cumsum, в соответствии с вышеуказанным groupby.В этом случае я хочу опустить первую группу из моего df.Мой желаемый df выглядит так:

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

Answer 1

Используйте GroupBy.transform с GroupBy.last для Series с тем же размером, что и оригинал DaatFrame, затем создайте столбец nw с помощью DataFrame.assign с 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