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