Los pandas obtienen los n primeros registros 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]})
que se parece a:
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 registros principales para cada identificación, como este:
id value
0 1 1
1 1 2
3 2 1
4 2 2
7 3 1
8 4 1
Puedo hacerlo numerando registros dentro del grupo después de groupby:
dfN = df.groupby('id').apply(lambda x:x['value'].reset_index()).reset_index()
que se parece a:
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
luego para la salida deseada:
dfN[dfN['level_1'] <= 1][['id', 'value']]
Producción:
id value
0 1 1
1 1 2
3 2 1
4 2 2
7 3 1
8 4 1
¿Pero existe un enfoque más eficaz/elegante para hacer esto? Y también existe un enfoque más elegante para los registros numéricos dentro de cada grupo (como la función de ventana SQL row_number() ).
Has probado
df.groupby('id').head(2)
Salida generada:
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 deba ordenar/ordenar antes, dependiendo de sus datos)
EDITAR: Como lo mencionó el autor de la pregunta, use
df.groupby('id').head(2).reset_index(drop=True)
para eliminar el MultiIndex y aplanar los resultados:
id value
0 1 1
1 1 2
2 2 1
3 2 2
4 3 1
5 4 1
Desde 0.14.1 , ahora puedes hacer nlargesty nsmallesten un 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
Existe una ligera rareza en el hecho de que también se incluye el índice original, pero esto podría ser realmente útil dependiendo de cuál era su índice original .
Si no está interesado en él, puede .reset_index(level=1, drop=True)deshacerse de él por completo.
(Nota: a partir de 0.17.1 también podrás hacer esto en DataFrameGroupBy, pero por ahora solo funciona con Seriesy SeriesGroupBy).
A veces, ordenar todos los datos requiere mucho tiempo. Podemos agrupar primero y hacer topk para cada grupo:
g = df.groupby(['id']).apply(lambda x: x.nlargest(topk,['value'])).reset_index(drop=True)
