Как суммировать nlargest() целые числа в GroupBy

Question

У меня есть dataframe так:

Index STNAME COUNTY COUNTY_POP 
     0  AL  0  100 
     1  AL  1  150 
     2  AL  3  200 
     3  AL  5  50 
    ... 
    15  CA  0  300 
    16  CA  1  200 
    17  CA  3  250 
    18  CA  4  350 

Я хочу суммировать три самых больших целых чисел от COUNTY_POP для каждого состояния. До сих пор у меня есть:

In[]: df.groupby(['STNAME'])['COUNTY_POP'].nlargest(3) 
    Out[]: 
    Index STNAME COUNTY COUNTY_POP 
     0  AL  0  100 
     1  AL  1  150 
     2  AL  3  200 
    ... 
    15  CA  0  300 
    17  CA  3  250 
    18  CA  4  350 

Однако при добавлении операции .sum() в коде выше, я получаю следующий результат.

In[]: df.groupby(['STNAME'])['COUNTY_POP'].nlargest(3).sum() 
    Out[]: 
    1350 

Я относительно новичок в Python и Pandas. Если бы кто-нибудь мог объяснить, что вызывает это и как его исправить, я бы очень признателен!

Answer 1

Это то, что вы хотите?

In [25]: df.groupby('STNAME')['COUNTY_POP'].agg(lambda x: x.nlargest(3).sum()) 
Out[25]: 
STNAME 
AL 450 
CA 900 
Name: COUNTY_POP, dtype: int64 

Answer 2

и ломтик сортировать почту перед отправкой на почтамт ... чуть быстрее

df.sort_values('COUNTY_POP').groupby('STNAME').COUNTY_POP \ 
    .apply(lambda x: x.values[-3:].sum()) 

STNAME 
AL 450 
CA 900 
Name: COUNTY_POP, dtype: int64