web-dev-qa-db-fra.com

Pandas groupby plus grande somme

J'essaie d'utiliser les fonctions groupby, nlargest et sum dans Pandas ensemble, mais j'ai du mal à le faire fonctionner.

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

Je veux utiliser groupby pour sélectionner par état, puis obtenir les 2 premiers comtés par population. Ensuite, utilisez uniquement les 2 premiers chiffres de la population du comté pour obtenir une somme pour cet état.

En fin de compte, je vais avoir une liste qui aura l'état et la population (de ses 2 premiers comtés).

Je peux faire fonctionner groupby et nlargest, mais obtenir la somme de nlargest(2) est un défi.

La ligne que j'ai en ce moment est simplement: df.groupby('State')['Population'].nlargest(2)

10
user7102752

Vous pouvez utiliser apply après avoir effectué le groupby:

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

Je pense que ce problème que vous rencontrez est que df.groupby('State')['Population'].nlargest(2) renverra un DataFrame, vous ne pouvez donc plus effectuer d'opérations au niveau du groupe. En général, si vous souhaitez effectuer plusieurs opérations dans un groupe, vous devrez utiliser apply/agg.

La sortie résultante:

State
Alabama    150
Wyoming    330

MODIFIER

Une approche légèrement plus propre, comme suggéré par @ cᴏʟᴅsᴘᴇᴇᴅ:

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

C'est cependant un peu plus lent que d'utiliser apply sur des DataFrames plus grands.

En utilisant la configuration suivante:

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

J'obtiens les horaires suivants:

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)

Le ralentissement des performances est potentiellement causé par le level kwarg dans sum effectuant un second groupby sous le capot.

23
root

En utilisant agg, la logique de regroupement ressemble à:

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

Il en résulte un autre objet de trame de données; que vous pouvez interroger pour trouver les États les plus peuplés, etc.

           Population
State
Alabama    150
Wyoming    330
4
aquaraga