Los pandas obtienen los n registros más altos dentro de cada grupo

Supongamos que tengo pandas DataFrame como este: 

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

Quiero obtener un nuevo DataFrame con los 2 mejores registros para cada id, así: 

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

Puedo hacerlo con la numeración de registros dentro de grupo tras grupo por: 

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

Pero, ¿hay un enfoque más eficaz/elegante para hacer esto? Y también hay un enfoque más elegante para los registros numéricos dentro de cada grupo (como SQL window function row_number()).

Author: smci, 2013-11-19
Source

2 answers

¿Intentaste df.groupby('id').head(2)

Resultado generado: 

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

(Tenga en cuenta que es posible que necesite ordenar/ordenar antes, dependiendo de sus datos)

EDITAR: Como mencionó el interrogador, use df.groupby('id').head(2).reset_index(drop=True) para eliminar el multindex y aplanar los 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
 110
Author: dorvak,
Warning: date(): Invalid date.timezone value 'Europe/Kyiv', we selected the timezone 'UTC' for now. in /var/www/agent_stack/data/www/ajaxhispano.com/template/agent.layouts/content.php on line 61
2013-11-19 11:01:41

Desde 0.14.1 , ahora puedes hacer nlargest y nsmallest en un objeto groupby: 

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

Hay una ligera rareza que se obtiene el índice original allí, así, pero esto podría ser realmente útil dependiendo de lo que su índice original era.

Si no estás interesado en él, puedes hacer .reset_index(level=1, drop=True) para deshacerte de él por completo.

(Nota: Desde 0.17.1 también podrás hacer esto en un DataFrameGroupBy, pero por ahora solo funciona con Series y SeriesGroupBy.)

 91
Author: LondonRob,
Warning: date(): Invalid date.timezone value 'Europe/Kyiv', we selected the timezone 'UTC' for now. in /var/www/agent_stack/data/www/ajaxhispano.com/template/agent.layouts/content.php on line 61
2015-09-04 12:14:13