Извлечение строк из кадра данных панд на основе условия

Question

У меня есть данные dataframe от pandas, и я хочу сохранить только те строки, в которых сумма «numb_people» для категории «class» составляет не менее 2.

Это, однако, приводит к ошибке индекса(индексы больше не совпадают):

data = data[data.groupby('class').sum()['numb_people'] > 2]

Как я могу сделать это так же просто?

Answer 1

Как сказал @rafaelc в комментарии:

idx = data.groupby('class').numb_people.transform('sum') > 2
print(data[idx])