Panda mendapatkan n teratas catatan dalam setiap kelompok

Misalkan saya memiliki panda DataFrame seperti ini:

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

Saya ingin mendapatkan DataFrame baru dengan 2 catatan teratas untuk setiap id, seperti ini:

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

Saya dapat melakukannya dengan penomoran catatan dalam grup demi grup dengan:

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

Tetapi apakah ada pendekatan yang lebih efektif / elegan untuk melakukan ini? Dan juga ada pendekatan yang lebih elegan untuk mencatat angka dalam setiap kelompok (seperti fungsi jendela SQL row_number () ).

Apakah kamu sudah mencoba? df.groupby('id').head(2)

Ouput dihasilkan:

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

(Perlu diingat bahwa Anda mungkin perlu memesan / mengurutkan sebelumnya, tergantung pada data Anda)

EDIT: Seperti yang disebutkan oleh penanya, gunakan df.groupby('id').head(2).reset_index(drop=True)untuk menghapus multindex dan meratakan hasilnya.

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

Sejak 0.14.1 , Anda sekarang dapat melakukan nlargestdan nsmallestpada groupbyobjek:

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

Ada keanehan sedikit bahwa Anda mendapatkan indeks asli di sana juga, tapi ini mungkin benar-benar berguna tergantung pada apa indeks asli Anda adalah .

Jika Anda tidak tertarik dengannya, Anda bisa melakukannya .reset_index(level=1, drop=True)untuk menghilangkannya sama sekali.

(Catatan: Dari 0.17.1 Anda juga dapat melakukan ini di DataFrameGroupBy, tetapi untuk saat ini hanya berfungsi dengan Seriesdan SeriesGroupBy.)

Terkadang mengurutkan seluruh data di depan sangat memakan waktu. Kita dapat mengelompokkannya terlebih dahulu dan melakukan topk untuk setiap grup:

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