Использование функций Pandas вместо циклов for для поиска трех верхних элементов в столбце

Question

У меня есть pandas фрейм данных, часть которого показана на этом рисунке (https://i.stack.imgur.com/VKJ8G.png). Поэтому мне нужно найти сумму населения (используя перепись 2010 года) 3 самых густонаселенных округов в каждом штате и найти три верхних штата с самой высокой суммой. Я новичок в pandas, поэтому я использовал циклы и условные выражения для выполнения sh задачи. Вот код, который я написал:

def answer_six():
    top1=["nil",0]
    top2=["nil",0]
    top3=["nil",0]
    for item in census_df['STNAME']:
        s_sum=census_df[census_df['STNAME'] == item].nlargest(3, ['CENSUS2010POP'])['CENSUS2010POP'].sum()
        if top1[1]==0:
            top1=[item,s_sum]
        if s_sum>top1[1]:
            top3=top2
            top2=top1
            top1=[item,s_sum]
        if s_sum<top1[1] and s_sum>top2[1]:
            top3=top2
            top2=[item,s_sum]
        if s_sum<top2[1] and s_sum>top3[1]:
            top3=[item,s_sum]
    ans=[top1[0],top2[0],top3[0]]
    return ans

Выполнение этого кода занимает приблизительно 15 секунд, чтобы получить ответ. Есть ли лучший способ использования функции pandas dataframe, такой как sum (), max () и т. Д. c.

Answer 1

Взято из обсуждения Code * , это будет работать:

census_df[census_df['SUMLEV'] == 50].groupby(
    'STNAME')['CENSUS2010POP'].apply(
    lambda x: x.nlargest(3).sum()).nlargest(
    3).index.values.tolist()