Pandas хороший подход к получению топ-n записей в каждой группе
Предположим, что у меня есть pandas DataFrame:
>>> df = pd.DataFrame({'id':[1,1,1,2,2,2,2,3,4],'value':[1,2,3,1,2,3,4,1,1]})
>>> df
id value
0 1 1
1 1 2
2 1 3
3 2 1
4 2 2
5 2 3
6 2 4
7 3 1
8 4 1
Я хочу получить новый DataFrame с топ-2 записями для каждого id, например:
id value
0 1 1
1 1 2
3 2 1
4 2 2
7 3 1
8 4 1
Я могу сделать это с нумерацией записей в группе после группы:
>>> dfN = df.groupby('id').apply(lambda x:x['value'].reset_index()).reset_index()
>>> dfN
id level_1 index value
0 1 0 0 1
1 1 1 1 2
2 1 2 2 3
3 2 0 3 1
4 2 1 4 2
5 2 2 5 3
6 2 3 6 4
7 3 0 7 1
8 4 0 8 1
>>> dfN[dfN['level_1'] <= 1][['id', 'value']]
id value
0 1 1
1 1 2
3 2 1
4 2 2
7 3 1
8 4 1
Но есть ли более эффективный/элегантный подход для этого? А также есть более элегантный подход к числовым записям в каждой группе (например, функция окна SQL row_number()).
Спасибо заранее.
Ответы
Ответ 1
Пробовал ли вы df.groupby('id').head(2)
Сгенерирован вывод:
>>> df.groupby('id').head(2)
id value
id
1 0 1 1
1 1 2
2 3 2 1
4 2 2
3 7 3 1
4 8 4 1
(Имейте в виду, что вам может потребоваться заказать/отсортировать до, в зависимости от ваших данных)
РЕДАКТИРОВАТЬ: Как упоминалось в вопросе, используйте df.groupby('id').head(2).reset_index(drop=True)
для удаления мультииндекса и сглаживания результатов.
>>> df.groupby('id').head(2).reset_index(drop=True)
id value
0 1 1
1 1 2
2 2 1
3 2 2
4 3 1
5 4 1
Ответ 2
Начиная с версии 0.14.1 теперь вы можете сделать nlargest
и nsmallest
в объекте groupby
:
In [23]: df.groupby('id')['value'].nlargest(2)
Out[23]:
id
1 2 3
1 2
2 6 4
5 3
3 7 1
4 8 1
dtype: int64
Там небольшая странность, что вы также получаете исходный индекс, но это может быть действительно полезно в зависимости от вашего исходного индекса.
Если вас это не интересует, вы можете сделать .reset_index(level=1, drop=True)
, чтобы полностью избавиться от него.
(Примечание: От 0.17.1 вы сможете сделать это и в DataFrameGroupBy, но пока он работает только с Series
и SeriesGroupBy
.)