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

Question

У меня есть кадр данных, показанный ниже. Он отсортирован так, что «POP» находится в порядке убывания по отношению к каждому «STATE». Теперь я хочу суммировать три самых больших значения «POP» в отношении каждого «СОСТОЯНИЯ», как мне это сделать?

import pandas as pd
d = [['X','q',123383],['X','w',43857349],['X','e',236657],['X','r',23574594],
        ['Y','t',547853],['Y','y',46282134],['Y','u',43857439],['Y','i',32654893],['Y','i',95678312]]
df = pd.DataFrame(d, columns = ['STATE','COUNTY','POP'])
df.sort_values(['STATE','POP'], ascending=[True, False]).set_index(['STATE','COUNTY'])

print(sorted_df)

# sorted_df:
                   POP
STATE COUNTY          
X     w       43857349
      r       23574594
      e         236657
      q         123383
Y     i       95678312
      y       46282134
      u       43857439
      i       32654893
      t         547853

Answer 1

Есть nlargest, который не требует предварительной сортировки:

df.groupby(['STATE']).POP.nlargest(3)

дает вам

STATE   
X      1    43857349
       3    23574594
       2      236657
Y      8    95678312
       5    46282134
       6    43857439
Name: POP, dtype: int64

и если вы заботитесь только о сумме:

df.groupby(['STATE']).POP.nlargest(3).sum(level=0)

дает:

STATE
X     67668600
Y    185817885
Name: POP, dtype: int64

Answer 2

Обязательно переназначьте после сортировки кадра данных (возможно, вы хотели назвать результат sorted_df).

Группируйте по уровню штата (или level=0, учитывая, что это мультииндекс штата и округа), затем применяйте лямбду, взяв голову в верхнюю тройку (учитывая, что данные были отсортированы в порядке убывания по населению для каждого состояния) и суммируем результат.

top_n = 3
df = df.sort_values(['STATE','POP'], ascending=[True, False]).set_index(['STATE','COUNTY'])
>>> df.groupby(level='STATE').apply(lambda x: x.head(top_n).sum())
             POP
STATE           
X       67668600  # w: 43857349 + r: 23574594 + e: 236657
Y      185817885  # i: 95678312 + y: 46282134 + u: 43857439