如何将Pandas中与failure相关性为正的列名加入列表？

问题 #

我是Pandas新手，想要把与failure列相关性大于0的DataFrame列名提取到列表中。我写了以下代码：

corr_matrix = df_train.corr()
corr_matrix["failure"].sort_values(ascending=False)

useful_features = []
for f in corr_matrix["failure"]:
    if f > 0:
        useful_features.append(df_train.columns)
print(useful_features)

但这段代码会把所有列名都添加到列表里，输出如下：

[Index(['id', 'product_code', 'loading', 'attribute_0', 'attribute_1',
       'attribute_2', 'attribute_3', 'measurement_0', 'measurement_1',
       'measurement_2', 'measurement_3', 'measurement_4', 'measurement_5',
       'measurement_6', 'measurement_7', 'measurement_8', 'measurement_9',
       'measurement_10', 'measurement_11', 'measurement_12', 'measurement_13',
       'measurement_14', 'measurement_15', 'measurement_16', 'measurement_17',
       'failure', 'kfold'],
. . .

我期望的输出是只包含符合条件的列名：

useful_features = ['failure','loading',...,'kfold']

另外，corr_matrix["failure"].sort_values(ascending=False)的输出结果如下：

failure           1.000000
loading           0.129089
measurement_17    0.033905
measurement_5     0.018079
measurement_8     0.017119
measurement_7     0.016787
measurement_2     0.015808
measurement_6     0.014791
measurement_0     0.009646
attribute_2       0.006337
measurement_14    0.006211
measurement_12    0.004398
measurement_3     0.003577
measurement_16    0.002237
kfold             0.000130
measurement_10   -0.001515
measurement_13   -0.001831
measurement_15   -0.003544
measurement_9    -0.003587
measurement_11   -0.004801
id               -0.007545
measurement_4    -0.010488
measurement_1    -0.010810
attribute_3      -0.019222
Name: failure, dtype: float64

请问怎么实现只添加符合条件的列名？用df_train.columns.values也会添加所有列名。

解决方案 #

问题出在循环中直接添加了整个列名集合df_train.columns，而非当前循环对应的单个列名。以下两种方法可以解决：

方法1：遍历索引与对应值 #

直接遍历corr_matrix["failure"]的索引（即列名）和相关性值，判断值大于0时将索引加入列表：

corr_matrix = df_train.corr()
# 获取排序后的failure相关性系列，也可以跳过排序直接处理
failure_corr = corr_matrix["failure"].sort_values(ascending=False)

useful_features = []
for col_name, corr_value in failure_corr.items():
    if corr_value > 0:
        useful_features.append(col_name)
print(useful_features)

方法2：布尔索引（Pandas简洁写法） #

利用Pandas布尔索引直接筛选符合条件的列名，无需手动循环：

corr_matrix = df_train.corr()
failure_corr = corr_matrix["failure"]
# 筛选相关性>0的列名并转成列表
useful_features = failure_corr[failure_corr > 0].index.tolist()

# 若需要保持排序后的顺序，先排序再筛选
useful_features = failure_corr.sort_values(ascending=False)[failure_corr > 0].index.tolist()

关键说明 #

• failure_corr.items()返回(列名, 相关性值)的元组，确保你能精准获取当前符合条件的单个列名。
• 布尔索引failure_corr > 0生成布尔系列，用它筛选原系列后，.index即为符合条件的列名，.tolist()可将其转换为Python列表。

内容的提问来源于stack exchange，提问作者Prathamesh Sawant