python - 使用 Pandas 函数而不是 for 循环来查找列中的前三个项目

Question

我有一个熊猫数据框，其中一部分显示在这张图片中（https://i.stack.imgur.com/VKJ8G.png）。所以我需要找到每个州人口最多的 3 个县的人口总和（使用 2010 年人口普查），并找到总和最高的前 3 个州。我是 pandas 的新手，所以我使用循环和条件来完成任务这是我写的代码：

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 秒才能得到答案。有没有更好的方法使用像 sum()、max() 等 pandas 数据框函数？

标签： pythonpandasdataframe

Answer 1

取自Code Review 讨论，这将起作用：

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