按n最大和分组的熊猫

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位的县。然后只使用前两个县人口数字来获得该州的总和。

最后，我将有一个列表，其中将包含州和人口(它的前两个县)。

我可以让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)

不过，这比在更大的DataFrames上使用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造成的。

Answer 2

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

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

这会产生另一个dataframe对象；您可以查询该对象以查找人口最多的状态，等等。

           Population
State
Alabama    150
Wyoming    330