Os pandas obtêm os mais altos n registros dentro de cada grupo

Suponha que eu tenha o DataFrame dos pandas assim:

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

Quero obter um novo DataFrame com os 2 principais registros para cada ID, assim:

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

Eu posso fazer isso com a numeração de registros dentro de grupo após grupo:

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

Mas existe uma abordagem mais eficaz / elegante para fazer isso? E também existe uma abordagem mais elegante para numerar registros dentro de cada grupo (como a função da janela SQL row_number () ).

Você tentou df.groupby('id').head(2)

Saída gerada:

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

Lembre-se de que talvez você precise solicitar / classificar antes, dependendo dos seus dados)

EDIT: Como mencionado pelo questionador, use df.groupby('id').head(2).reset_index(drop=True)para remover o multíndice e achatar os resultados.

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

Desde 0.14.1 , agora você pode fazer nlargeste nsmallestem um groupbyobjeto:

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

Há uma ligeira estranheza que você obtenha o índice original, em lá também, mas isso pode ser realmente útil dependendo do que o seu índice de originais foi .

Se você não estiver interessado, pode fazer isso .reset_index(level=1, drop=True)para se livrar completamente.

(Nota: A partir da 0.17.1, você também poderá fazer isso em um DataFrameGroupBy, mas por enquanto ele só funciona com Seriese SeriesGroupBy.)

Às vezes, classificar todos os dados com antecedência leva muito tempo. Podemos agrupar primeiro e fazer topk para cada grupo:

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