Как собрать два самых больших значения на группу в пандах?

Question

Я проходил по этой ссылке: Вернул N самых больших значений в каждой группе, используя панд

, и нашел несколько способов найти значения topN в группе.

ОднакоЯ предпочитаю метод словаря с функцией agg и хотел бы знать, возможно ли получить эквивалент метода словаря для следующей задачи?

import numpy as np
import pandas as pd

df = pd.DataFrame({'A': [1,   1,   1,   2,   2],
                   'B': [1,   1,   2,   2,   1],
                   'C': [10,  20,  30,  40,  50],
                   'D': ['X', 'Y', 'X', 'Y', 'Y']})
print(df)
   A  B   C  D
0  1  1  10  X
1  1  1  20  Y
2  1  2  30  X
3  2  2  40  Y
4  2  1  50  Y

Я могу сделать это:

df1 = df.groupby(['A'])['C'].nlargest(2).droplevel(-1).reset_index()
print(df1)
   A   C
0  1  30
1  1  20
2  2  50
3  2  40

# also this
df1 = df.sort_values('C', ascending=False).groupby('A', sort=False).head(2)
print(df1)

# also this
df.set_index('C').groupby('A')['B'].nlargest(2).reset_index()

Обязательно

df.groupby('A',as_index=False).agg(
    {'C': lambda ser: ser.nlargest(2)  # something like this
    })

Можно ли здесь использовать словарь?

Answer 1

Если вы хотите получить словарь, подобный A: 2 верхних значения из C , вы можете запустить:

df.groupby(['A'])['C'].apply(lambda x:
    x.nlargest(2).tolist()).to_dict()

Для вашего DataFrame результат будет:

{1: [30, 20], 2: [50, 40]}