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

Предположим, у меня есть 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 ()).

Источник

6 ответов

Решение

Ты пробовал 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
Источник

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

Источник

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

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

Чтобы получить первые N строк каждой группы, можно использовать другой способ:groupby().nth[:N]. Результат этого вызова такой же, какgroupby().head(N). Например, для первых двух строк для каждого идентификатора вызовите:

      N = 2
df1 = df.groupby('id', as_index=False).nth[:N]

Чтобы получить наибольшее значение N каждой группы, я предлагаю два подхода.

  1. Сначала отсортируйте по «id» и «value» (не забудьте отсортировать «id» в порядке возрастания, а «value» в порядке убывания, используяascendingпараметр соответствующим образом), а затем вызовитеgroupby().nth[].

            N = 2
df1 = df.sort_values(by=['id', 'value'], ascending=[True, False])
df1 = df1.groupby('id', as_index=False).nth[:N]

  2. Другой подход заключается в ранжировании значений каждой группы и фильтрации с использованием этих рангов.

            # for the entire rows
N = 2
msk = df.groupby('id')['value'].rank(method='first', ascending=False) <= N
df1 = df[msk]

# for specific column rows
df1 = df.loc[msk, 'value']

Оба они намного быстрее, чемgroupby().apply()иgroupby().nlargest()звонки, как предложено в других ответах здесь ( 1 , 2, 3 ). В образце из 100 тыс. строк и 8000 групп%timeitтест показал, что это было в 24-150 раз быстрее, чем эти решения.

Кроме того, вместо нарезки вы также можете передать список/кортеж/диапазон в.nth()вызов:

      df.groupby('id', as_index=False).nth([0,1])

# doesn't even have to be consecutive
# the following returns 1st and 3rd row of each id
df.groupby('id', as_index=False).nth([0,2])
Источник 
      df.groupby('id').apply(lambda x : x.sort_values(by = 'value', ascending = False).head(2).reset_index(drop = True))
  • Здесь значения сортировки по возрастанию false дают аналогичные nlargest, а True - аналогичные nsmallest.
  • Значение внутри заголовка совпадает со значением, которое мы даем внутри nlargest, чтобы получить количество значений, отображаемых для каждой группы.
  • reset_index является необязательным и необязательным.
Источник

Это работает для повторяющихся значений

Если у вас есть повторяющиеся значения в первых n значениях и вам нужны только уникальные значения, вы можете сделать это следующим образом:

      import pandas as pd

ifile = "https://raw.githubusercontent.com/bhishanpdl/Shared/master/data/twitter_employee.tsv"
df = pd.read_csv(ifile,delimiter='\t')
print(df.query("department == 'Audit'")[['id','first_name','last_name','department','salary']])

    id first_name last_name department  salary
24  12   Shandler      Bing      Audit  110000
25  14      Jason       Tom      Audit  100000
26  16     Celine    Anston      Audit  100000
27  15    Michale   Jackson      Audit   70000

If we do not remove duplicates, for the audit department we get top 3 salaries as 110k,100k and 100k.
If we want to have not-duplicated salaries per each department, we can do this:

(df.groupby('department')['salary']
 .apply(lambda ser: ser.drop_duplicates().nlargest(3))
 .droplevel(level=1)
 .sort_index()
 .reset_index()
)

This gives

department  salary
0   Audit   110000
1   Audit   100000
2   Audit   70000
3   Management  250000
4   Management  200000
5   Management  150000
6   Sales   220000
7   Sales   200000
8   Sales   150000
Источник
Другие вопросы по тегам