Pandas erhalten in jeder Gruppe die höchsten n Datensätze

Angenommen, ich habe Pandas DataFrame wie folgt:

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

Ich möchte einen neuen DataFrame mit den beiden besten Datensätzen für jede ID erhalten:

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

Ich kann es mit der Nummerierung von Datensätzen innerhalb einer Gruppe nach der anderen tun, indem ich:

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

Aber gibt es dafür einen effektiveren / eleganteren Ansatz? Außerdem gibt es einen eleganteren Ansatz für die Nummerierung von Datensätzen in jeder Gruppe (wie die SQL-Fensterfunktion row_number () ).

Hast du versucht df.groupby('id').head(2)

Ausgabe generiert:

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

(Beachten Sie, dass Sie je nach Ihren Daten möglicherweise vorher bestellen / sortieren müssen.)

BEARBEITEN: Verwenden Sie df.groupby('id').head(2).reset_index(drop=True), wie vom Fragesteller erwähnt, um den Multindex zu entfernen und die Ergebnisse zu reduzieren.

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

Seit 0.14.1 können Sie jetzt nlargestund nsmallestauf einem groupbyObjekt:

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

Es ist etwas seltsam, dass Sie dort auch den Originalindex erhalten, aber dies kann sehr nützlich sein, je nachdem, was Ihr Originalindex war .

Wenn Sie nicht daran interessiert sind, können Sie .reset_index(level=1, drop=True)es ganz loswerden.

(Hinweis: Ab 0.17.1 können Sie dies auch auf einem DataFrameGroupBy tun, aber im Moment funktioniert es nur mit Seriesund SeriesGroupBy.)

Manchmal ist das Sortieren der gesamten Daten sehr zeitaufwändig. Wir können zuerst gruppieren und für jede Gruppe Topk machen:

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