Les pandas obtiennent les n enregistrements les plus élevés dans chaque groupe

164

Supposons que j'ai des pandas DataFrame comme ceci:

>>> df = pd.DataFrame({'id':[1,1,1,2,2,2,2,3,4],'value':[1,2,3,1,2,3,4,1,1]})
>>> df
   id  value
0   1      1
1   1      2
2   1      3
3   2      1
4   2      2
5   2      3
6   2      4
7   3      1
8   4      1

Je souhaite obtenir un nouveau DataFrame avec les 2 meilleurs enregistrements pour chaque identifiant, comme ceci:

   id  value
0   1      1
1   1      2
3   2      1
4   2      2
7   3      1
8   4      1

Je peux le faire avec la numérotation des enregistrements dans un groupe après un groupe en:

>>> dfN = df.groupby('id').apply(lambda x:x['value'].reset_index()).reset_index()
>>> dfN
   id  level_1  index  value
0   1        0      0      1
1   1        1      1      2
2   1        2      2      3
3   2        0      3      1
4   2        1      4      2
5   2        2      5      3
6   2        3      6      4
7   3        0      7      1
8   4        0      8      1
>>> dfN[dfN['level_1'] <= 1][['id', 'value']]
   id  value
0   1      1
1   1      2
3   2      1
4   2      2
7   3      1
8   4      1

Mais y a-t-il une approche plus efficace / élégante pour faire cela? Et y a-t-il aussi une approche plus élégante des enregistrements numériques dans chaque groupe (comme la fonction de fenêtre SQL row_number () ).

Roman Pekar
la source
Une duplication possible de la trame
ssoler
1
"top-n" ne signifie pas "les n rangées les plus élevées / les premières / têtes", comme vous le recherchez! Cela signifie "les n lignes avec les plus grandes valeurs".
smci

Réponses:

183

As-tu essayé df.groupby('id').head(2)

Ouput généré:

>>> df.groupby('id').head(2)
       id  value
id             
1  0   1      1
   1   1      2 
2  3   2      1
   4   2      2
3  7   3      1
4  8   4      1

(Gardez à l'esprit que vous devrez peut-être commander / trier avant, en fonction de vos données)

EDIT: Comme mentionné par l'interrogateur, utilisez df.groupby('id').head(2).reset_index(drop=True)pour supprimer le multindex et aplatir les résultats.

>>> df.groupby('id').head(2).reset_index(drop=True)
    id  value
0   1      1
1   1      2
2   2      1
3   2      2
4   3      1
5   4      1
Dorvak
la source
1
Oui, je pense que c'est ça. Oublié cela en quelque sorte. Connaissez-vous un bon moyen de numéroter des enregistrements au sein d'un groupe?
Roman Pekar
4
Pour obtenir la sortie dont j'ai besoin, j'ai également ajouté.reset_index(drop=True)
Roman Pekar
1
github.com/pydata/pandas/pull/5510 vient d'être fusionné; sera en 0.13, nouvelle méthode pour faire exactement cela appelée cumcount(numéroter les enregistrements dans chaque groupe)
Jeff
1
@Jeff bonne nouvelle. J'aurais aimé avoir plus de temps pour contribuer à Pandas :(
Roman Pekar
3
Pour rendre @dorvak sa réponse plus complète, si vous voulez les 2 plus petites valeurs, idfaites-le df.sort_values(['id', 'value'], axis=0).groupby('id').head(2). Autre exemple, la plus grande valeur par idest donnée par df.sort_values(['id', 'value'], axis=0).groupby('id').tail(1).
Elmex80s
133

Depuis la 0.14.1 , vous pouvez maintenant faire nlargestet nsmallestsur un groupbyobjet:

In [23]: df.groupby('id')['value'].nlargest(2)
Out[23]: 
id   
1   2    3
    1    2
2   6    4
    5    3
3   7    1
4   8    1
dtype: int64

Il y a une légère bizarreries que vous obtenez l'index original là - dedans aussi, mais cela pourrait être vraiment utile en fonction de ce que votre index d' origine était .

Si cela ne vous intéresse pas, vous pouvez vous .reset_index(level=1, drop=True)en débarrasser complètement.

(Remarque: à partir de la version 0.17.1, vous pourrez également le faire sur un DataFrameGroupBy, mais pour l'instant, cela ne fonctionne qu'avec Serieset SeriesGroupBy.)

LondresRob
la source
Y a-t-il un moyen d'y arriver unique_limit(n)? Comme je veux les n premières valeurs uniques? Si je le demande, nlargestil
triera
2
Cela ne fonctionne pas pour les cas où vous effectuez un agrégat sur le groupby? Par exemple, df.groupby([pd.Grouper(freq='M'), 'A'])['B'].count().nlargest(5, 'B') cela renvoie simplement le top 5 général de la série entière, pas par chaque groupe
géominded
L'affirmation selon laquelle cela est désormais également possible sur DataFrameGroupBys semble être fausse, la demande d'extraction liée semble s'ajouter uniquement aux s nlargestsimples DataFrame. Ce qui est plutôt malheureux, car que se passe-t-il si vous souhaitez sélectionner plus d'une colonne?
oulenz
7

Parfois, le tri de toutes les données à l'avance prend beaucoup de temps. Nous pouvons grouper d'abord et faire topk pour chaque groupe:

g = df.groupby(['id']).apply(lambda x: x.nlargest(topk,['value'])).reset_index(drop=True)
Chaffee Chen
la source