A la hora de seleccionar las características de un data frame, es posible que nos encontremos con la necesidad de seleccionar o eliminar características del data frame y que el nombre de esas características tenga un determinado patrón. Esta labor la podemos realizar mediante selección de elementos en listas; en esta entrada del blog vamos a tener tres tipos de selecciones:

  1. Seleccionar o eliminar aquellas variables que empiezan por un determinado prefijo.
  2. Seleccionar o eliminar aquellas variables que contienen una cadena de caracteres.
  3. Seleccionar o eliminar aquellas variables que finalizan con un sufijo.

Para ilustrar este trabajo, generamos un data frame con datos aleatorios y 10 columnas:

import numpy as np
import pandas as pd

df = pd.DataFrame(np.random.randint(0, 100, size=(100, 10)),
                  columns=['A1', 'A2', 'A3', 'B1', 'B2', 'B3', 'C1', 'C2', 'C3', 'DA'])

El primero de los filtros a realizar es identificar qué variables de nuestro data frame contienen el string 'A':

# Filtro 1: Columnas que tienen una 'A'
filtro1 = [col for col in df.columns if col.find('A') >= 0]

# Eliminar
df1_drop = df.drop(columns=filtro1)

# Seleccionar
df1_keep = df[filtro1]

Siempre vamos a hacer el mismo proceso: las características de nuestro data frame irán en una lista; después recorremos la lista y seleccionamos aquellos donde el método .find('A') sea mayor o igual a 0; con esto hemos creado una sublista con aquellas características que tienen el string 'A'. Mediante .drop(columns=) eliminamos del data frame los elementos contenidos en una lista.

El siguiente filtro se queda con las variables que empiezan por 'A' o bien las elimina:

# Filtro 2: Columnas que empiezan por una 'A'
filtro2 = [col for col in df.columns if col.startswith('A')]

# Eliminar
df2_drop = df.drop(columns=filtro2)

# Seleccionar
df2_keep = df[filtro2]

En este caso empleamos el método .startswith() para seleccionar sólo las variables que empiecen por A. Y por último, realizamos un filtro en base a un sufijo, si terminan por '1':

# Filtro 3: Columnas que finalizan por un '1'
filtro3 = [col for col in df.columns if col.endswith('1')]

# Eliminar
df3_drop = df.drop(columns=filtro3)

# Seleccionar
df3_keep = df[filtro3]

Desde mi prisma, es un código muy útil ya que tengo la costumbre de incluir en los modelos sólo variables con un determinado prefijo y, como perro viejo, estoy haciendo lo mismo con mis modelos de Machine Learning; así siempre tengo claro qué variables y en qué paso las ha incluido el proceso, por muy automático que sea. Saludos.