Панди отримують найвищі n записів у кожній групі


162

Припустимо, у мене є такі панди 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 () ).


Можливий дублікат фрейму
ssoler

1
"Топ-п" не означає "п. найвищих / перших / головних рядків", як ви шукаєте! Це означає "n рядів з найбільшими значеннями".
smci

Відповіді:


183

Ви спробували 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

(Майте на увазі, що вам може знадобитися замовити / сортувати раніше, залежно від ваших даних)

EDIT: Як згадував 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

1
Так, я думаю, що це все. Це якимось чином не помітив. Чи знаєте ви хороший спосіб нумерації записів у групі?
Роман Пекар

4
Щоб отримати потрібний результат, я також додав.reset_index(drop=True)
Роман Пекар

1
github.com/pydata/pandas/pull/5510 щойно було об'єднано; буде в 0.13, новий метод зробити саме так називається cumcount(пронумеруйте записи в кожній групі)
Jeff

1
@Jeff хороші новини. Я б хотів, щоб у мене було більше часу, щоб сприяти «Пандам» :(
Роман Пекар

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

132

З 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.)


Є спосіб дістатися unique_limit(n)? Наче я хочу перших n унікальних цінностей? Якщо я попрошу, я nlargestвідсортую весь df, який може бути дорогим
citynorman

2
Це не працює у випадках, коли ви робите сукупність у групі? Наприклад, df.groupby([pd.Grouper(freq='M'), 'A'])['B'].count().nlargest(5, 'B') це просто повертає загальні топ-5 у всій серії, а не в кожній групі
задумав

Заява про те, що тепер це також можливо на DataFrameGroupBys, видається помилковим, схожий запит на потяг додається лише nlargestдо простого DataFrames. Що досить прикро, адже що робити, якщо ви хочете вибрати більше одного стовпця?
oulenz

7

Іноді сортування цілих даних попереду дуже забирає багато часу. Ми можемо групуватися першими і робити топк для кожної групи:

g = df.groupby(['id']).apply(lambda x: x.nlargest(topk,['value'])).reset_index(drop=True)
Використовуючи наш веб-сайт, ви визнаєте, що прочитали та зрозуміли наші Політику щодо файлів cookie та Політику конфіденційності.
Licensed under cc by-sa 3.0 with attribution required.