Панды получают лучшие n записей в каждой группе

164

Предположим, у меня есть 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 с двумя верхними записями для каждого идентификатора, например:

   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
"top-n" не означает "n topmost / first / head row", как вы ищете! Это означает «n строк с самыми большими значениями».
smci

Ответы:

183

Ты пробовал df.groupby('id').head(2)

Ouput генерируется:

>>> 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
dorvak
источник
1
Да, я думаю это все. Это как-то упустили. Знаете ли вы хороший способ нумерации записей в группе?
Роман Пекарь
4
Чтобы получить выходной мне нужно, я также добавил.reset_index(drop=True)
Роман Пекар
1
github.com/pydata/pandas/pull/5510 был только что объединен; будет в 0.13, новый метод, чтобы сделать именно это называется cumcount(нумерация записей в каждой группе)
Джефф
1
@ Джефф хорошие новости. Хотелось бы, чтобы у меня было больше времени, чтобы внести свой вклад в Панды :(
Роман Пекар
3
Чтобы сделать @dorvak его ответ более полным, если вы хотите получить 2 наименьших значения за, idтогда сделайте df.sort_values(['id', 'value'], axis=0).groupby('id').head(2). Еще один пример, наибольшее значение за idдается df.sort_values(['id', 'value'], axis=0).groupby('id').tail(1).
Elmex80s
133

Начиная с 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.)

LondonRob
источник
Есть ли способ получить unique_limit(n)? Как я хочу первые n уникальных значений? Если я попрошу об nlargestэтом, я расскажу весь df, который может быть дорогим
citynorman
2
Это не работает для случаев, когда вы делаете агрегирование по группам? Например, df.groupby([pd.Grouper(freq='M'), 'A'])['B'].count().nlargest(5, 'B') это просто возвращает общий топ-5 во всей серии, а не по каждой группе
geominded
Утверждение о том, что это теперь также возможно для DataFrameGroupBys, кажется ложным, связанный запрос извлечения, кажется, добавляет только nlargestк простым DataFrames. Что довольно прискорбно, потому что, если вы хотите выбрать более одного столбца?
oulenz
7

Иногда сортировка целых данных занимает очень много времени. Мы можем сначала сгруппироваться и сделать topk для каждой группы:

g = df.groupby(['id']).apply(lambda x: x.nlargest(topk,['value'])).reset_index(drop=True)
Чаффи Чен
источник