Pandas: estadísticas de subconjuntos de datos en DataFrame

A partir de un conjunto de datos en memoria puede ser interesante obtener estadísticas de subconjuntos de datos. Información que puede ser de gran utilidad para detectar patrones. Por ejemplo, si tenemos los registros de ventas con el día de la semana puede ser interesante obtener los valores medios y totales en función del día. En Python esto se puede hacer gracias al método groupby() de los DataFrame Pandas.

Conjunto de dato de ejemplo

Para ver cómo sacar provecho del método groupby() de Pandas y poder obtener estadísticas de subconjuntos de datos en primer lugar es necesario disponer de un conjunto de datos de ejemplo, como el que se puede crear con el siguiente código.

import pandas as pd

df = pd.DataFrame({'C1': ["x", "x", "x", "x", "x", "y", "y", "y", "y"],
                   'C2': ["a", "b", "a", "a", "b", "a", "b", "b", "b"],
                   'C3': [12, 2, 7, 11, 3, 6, 9, 11, 7],
                   'C4': [4, 7, 1, 3, 7, 3, 14, 3, 8]})

Un conjunto de datos que se ha creado con dos características categóricas ( C1 y C2) y dos numéricas ( C3 y C4). Pudiéndose agrupar los registros en base a las categóricas. Por ejemplo, se puede insertar calcular la media de las variables numéricas en función de los valores que tienen los registros de C1.

Media de los valores los valores de las características

Para obtener la media de las columnas C3 y C4 para cada valor de C1 solamente hay que agrupar por esta mediante el uso de groupby() y emplear la función del estadístico deseado. Solamente hay que pasar como parámetro de groupby() el nombre de la columnas que se desea usar para agrupar los valores.

df.groupby('C1').mean()

      C3   C4
C1           
x   7.00  4.4
y   8.25  7.0

Así, de una manera sencilla se puede ver que la media de C3 para los valores con x en C1 es 7 y 4.4 para el caso de los valores de C3. La agrupación también se puede hacer en base a los valores de más de una columnas, para lo que es necesario pasar como parámetro del método groupby() la lista de columnas. De este modo, para obtener la media en base a los valores de C1 y C2 solamente hay que usar el siguiente código.

df.groupby(['C1', 'C2']).mean()

         C3        C4
C1 C2                
x  a   10.0  2.666667
   b    2.5  7.000000
y  a    6.0  3.000000
   b    9.0  8.333333

Contar el número de registros por agrupación

Por otro lado, si lo que se desea es obtener el número de registros por agrupación solamente hay que cambiar el método mean() por size().

df.groupby(['C1', 'C2']).size()

C1  C2
x   a     3
    b     2
y   a     1
    b     3
dtype: int64

En este caso se puede ver que la columna de resultados no tiene nombre, algo que se puede solucionar con el método reset_index().

df.groupby(['C1', 'C2']).size().reset_index(name='counts')

  C1 C2  counts
0  x  a       3
1  x  b       2
2  y  a       1
3  y  b       3

Obtener múltiples estadísticos

Lo que se ha visto anteriormente está bien, pero en muchos casos puede que se desee obtener más de un estadístico para cada grupo. Por ejemplo, la media, el mínimo, el máximo y el número de registros. Para estos casos se cuenta con el método agg() gracias al cual se pueden aplicar más de una operación sobre cada elemento, simplemente pasaron el nombre de estas como parámetro del método. Así para conseguir el ejemplo propuesto simplemente hay que escribir.

df.groupby(['C1', 'C2']).agg(['mean', 'max', 'min', 'count'])

         C3                      C4              
       mean max min count      mean max min count
C1 C2                                            
x  a   10.0  12   7     3  2.666667   4   1     3
   b    2.5   3   2     2  7.000000   7   7     2
y  a    6.0   6   6     1  3.000000   3   3     1
   b    9.0  11   7     3  8.333333  14   3     3

En este caso obtener el número de registros de todas las columnas puede ser algo repetitivo, además puede que deseemos aplicar diferentes estadísticos para cada una de las columnas. Para lo que también existen una solución, usar un diccionario con las funciones que se desean aplicar a cada una de las columnas. Pudiendo, por ejemplo, la media y el mínimo para C3 y la media, el máximo y el número de registros para C4.

df.groupby(['C1', 'C2']).agg({'C3': ['mean', 'min'],
                              'C4': ['mean', 'max', 'count']})

         C3            C4          
       mean min      mean max count
C1 C2                              
x  a   10.0   7  2.666667   4     3
   b    2.5   2  7.000000   7     2
y  a    6.0   6  3.000000   3     1
   b    9.0   7  8.333333  14     3

Cambiar en nombre de las columnas

Finalmente, el resultado que se ha obtenido tiene las etiquetas de las columnas en niveles. En muchos casos puede ser más interesante contar con diferente en lugar de niveles. Algo que se puede solucionar concatenados los nombres de los niveles y asignando el resultado a los nombres de las columnas. Lo que se puede conseguir fácilmente utilizando el método map(), un método que se ha visto al hablar de programación funcional.

df = df.groupby(['C1', 'C2']).agg({'C3': ['mean', 'min'],
                                   'C4': ['mean', 'max', 'count']})

df.columns = df.columns.map('_'.join).str.strip('_')
df

       C3_mean  C3_min   C4_mean  C4_max  C4_count
C1 C2                                             
x  a      10.0       7  2.666667       4         3
   b       2.5       2  7.000000       7         2
y  a       6.0       6  3.000000       3         1
   b       9.0       7  8.333333      14         3

Conclusiones

En esta entrada se ha visto cómo se puede obtener estadísticas de subconjuntos de datos en DataFrame Pandas gracias al método groupby(). Un método con el que se puede conseguir grandes resultados con poco trabajo. Consiguiendo además un código que es bastante legible.