Máximos de un subconjunto en un dataframe pandas

Una tarea que se puede hacer fácilmente en SQL agrupar una tabla por una columna, seleccionar los registros que tengan el valor máximo en otra columna, junto al resto de campos de los registros. Es decir, buscar los valores máximos de un subconjunto. Una tarea que también se puede hacer en pandas con un poco de trabajo.

Planteamiento del problema

Si tenemos una tabla como la siguiente.

Se desea buscar los registros máximos de la columna “C3” para cada uno de los valores diferente de la columna “C1”. Es decir, cuando para “x” se busca el registro con valor 12 en “C3”, para “y” con valor 9 y para “z” 11. Además, se quiere obtener los valores asociados en la columna “C2” o cualquier otra que pueda tener la tabla.

Uso de gropuBy y agg para filtrar en subconjunto de un dataframe.

Una posible solución al problema es usar el comando gropuBy junto a agg para identificar la posición de los máximos en el dataframe. En el caso de conocer los índices de los registros, estos se pueden utilizar para como filtros de selección.

Como parámetro del método agg se puede utilizar el nombre de una función o, más genéricamente, una función lambda. Lo que le llega a esta función son los registros de la columna que se corresponden con la agrupación. Así si se indica que el método se aplique sobre la columna “C3” se puede obtener los índices en los que se encuentran los máximos con idmax()o los mínimos con idmin(). Es decir, la posición de los máximos se puede obtener mediante el siguiente código.

import pandas as pd

df = pd.DataFrame({'C1': ["x", "x", "x", "y", "y", "y", "z"],
                   'C2': ["a", "b", "c", "d", "e", "f", "g"],
                   'C3': [12,  2,   7,   3,   6,   9,   11]})

df.groupby('C1').agg(max_ = ('C3', lambda data: data.idxmax()))

    max_
C1      
x      0
y      5
z      6

Ahora simplemente se puede utilizar esta columna para filtrar los resultados que se desean extraer.

df.iloc[df.groupby('C1').agg(max_ = ('C3', lambda data: data.idxmax())).max_]

  C1 C2  C3
0  x  a  12
5  y  f   9
6  z  g  11

Siendo posible seleccionar solamente algunas de las columnas del dataframe.

Conclusiones

En esta entrada se ha visto cómo los valores máximos de un subconjunto en un dataframe pandas. Aunque el método propuesto también se podría utilizar para seleccionar los mínimos. Un truco que podría de utilidad en algunas situaciones.

Imagen de nextvoyage en Pixabay