Pandas groupby nlargest sum

Posted 2023-03-11

【中文标题】Pandas groupby nlargest sum

【英文标题】：

【发布时间】：2017-03-16 09:30:41

【问题描述】：

我正在尝试在 Pandas 中同时使用 groupby、nlargest 和 sum 函数，但无法使其正常工作。

State    County    Population
Alabama  a         100
Alabama  b         50
Alabama  c         40
Alabama  d         5
Alabama  e         1
...
Wyoming  a.51      180
Wyoming  b.51      150
Wyoming  c.51      56
Wyoming  d.51      5

我想使用groupby 按州选择，然后按人口获得前 2 个县。然后仅使用前 2 个县的人口数来获得该州的总和。

最后，我将列出一个包含州和人口（前 2 个县）的列表。

我可以让groupby 和nlargest 工作，但是获得nlargest(2) 的总和是一个挑战。

我现在的行很简单：df.groupby('State')['Population'].nlargest(2)

【参考方案1】：

你可以在执行groupby之后使用apply：

df.groupby('State')['Population'].apply(lambda grp: grp.nlargest(2).sum())

我认为您遇到的这个问题是 df.groupby('State')['Population'].nlargest(2) 将返回一个 DataFrame，因此您不能再进行组级操作。一般来说，如果要在一个组中执行多个操作，则需要使用apply/agg。

结果输出：

State
Alabama    150
Wyoming    330

编辑

@cᴏʟᴅsᴘᴇᴇᴅ 建议的更简洁的方法：

df.groupby('State')['Population'].nlargest(2).sum(level=0)

这比在较大的 DataFrame 上使用 apply 稍慢。

使用以下设置：

import numpy as np
import pandas as pd
from string import ascii_letters

n = 10**6
df = pd.DataFrame('A': np.random.choice(list(ascii_letters), size=n),
                   'B': np.random.randint(10**7, size=n))

我得到以下时间：

In [3]: %timeit df.groupby('A')['B'].apply(lambda grp: grp.nlargest(2).sum())
103 ms ± 1.08 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

In [4]: %timeit df.groupby('A')['B'].nlargest(2).sum(level=0)
147 ms ± 3.38 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

较慢的性能可能是由于sum 中的level kwarg 在后台执行第二个groupby 造成的。

【讨论】：

你应该通过摆脱应用来解决这个问题，使用df.groupby('State')['Population'].nlargest(2).sum(level=0)

@cᴏʟᴅsᴘᴇᴇᴅ：当数据框大小很大时，您提出的.nlargest(2).sum(level=0) 解决方案实际上比使用apply 慢。 sum 中的 level kwarg 在后台执行第二个 groupby 操作，我猜这就是额外开销的来源。

这太令人惊讶了。所以一个 groupby + apply 胜过两个 groupbys。学到了新东西，干杯！

如果countys 不是唯一的，您的解决方案是否有效？想象一下有 2 行，例如：Alabama;e;2 和 Alabama;e;39？是否会考虑使用e 而不是c？

【参考方案2】：

使用agg，分组逻辑如下：

df.groupby('State').agg('Population': lambda x: x.nlargest(2).sum() )

这会产生另一个数据框对象；您可以查询以找到人口最多的州等。

           Population
State
Alabama    150
Wyoming    330