Cómo seleccionar columnas que están altamente correlacionadas con una columna específica en un marco de datos: python, pandas, correlación, tabular

Tengo un marco de datos que tiene más de 100 columnas,Con lo cual estoy tratando de construir un modelo. En este caso, una columna (A) en este marco de datos se considera como una respuesta y todas las demás columnas (B, C, D, etc.) son predictores. Así que estoy tratando de seleccionar todas las columnas que están correlacionadas con la columna A en función del factor de correlación (digamos> 0.2). Ya generé un mapa de calor con todos los factores de correlación entre cada par de columnas. Pero, ¿puedo tener un método rápido en pandas para obtener todas las columnas con un factor de correlación de más de 0.2 (que, por supuesto, ajustaré si es necesario) a la columna A? ¡Gracias por adelantado!

Respuestas

Utilizar el DataFrame para calcular la correlación, luego corte las columnas según su condición de corte con una máscara booleana.

import pandas as pd
df = pd.DataFrame({"A": [1,2,3,4,5,6,7,8,9,10],
"B": [1,2,4,3,5,7,6,8,10,11],
"C": [15,-1,17,-10,-10,-13,-99,-101,0,0],
"D": [0,10,0,0,-10,0,0,-10,0,10]} )

df.loc[:, df.corr()["A"] > 0.2]

A   B
0   1   1
1   2   2
2   3   4
3   4   3
4   5   5
5   6   7
6   7   6
7   8   8
8   9   10
9   10  11