Pandas groupby nlargest sum

Question

我试图在 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)

Answer 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)

较慢的性能可能通过引起level在kwarg sum执行第二groupby罩下。

Answer 2

使用agg ，分组逻辑如下所示：

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

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

           Population
State
Alabama    150
Wyoming    330