Jak zsumować nlargest() całkowitymi w GroupBy

Question

Mam dataframe tak:

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 

Chcę podsumować trzy największe liczby całkowite od COUNTY_POP dla każdego stanu. Do tej pory mam:

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 

Jednak po dodaniu .sum() operację powyższym kodzie, otrzymuję następujący wynik.

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

Jestem stosunkowo nowy w Pythonie i Pandach. Jeśli ktokolwiek mógłby wyjaśnić, co to powoduje i jak to poprawić, byłbym bardzo wdzięczny!

Answer 1

Czy tego chcesz?

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

presort i kawałek ... ciut szybciej

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