Scikit Learn 基于树的特征选择保留列名

Question

我想选择基于树的特征。我的数据集大约有 30 列，完成后大约有 5 列。这对我来说很棒，我遇到的问题是我得到的 5 列数据集不保留列的名称，我无法识别它们。

import pandas as pd
from sklearn.ensemble import ExtraTreesClassifier
from sklearn.feature_selection import SelectFromModel
data = pd.read_csv(file)
X = data.drop('target', 1)
y = data['target']
X.shape                        #(100000, 30)
clf = ExtraTreesClassifier()
clf = clf.fit(X, y)
clf.feature_importances_  
model = SelectFromModel(clf, prefit=True)
X_new = model.transform(X)
X_new.shape                    #(100000, 5)

有人可以帮我吗？

Answer 1

现在，当我更确定答案时，请尝试以下操作：

mask = model.get_support(indices=False) # this will return boolean mask for the columns X_new = X.loc[:, mask] # the sliced dataframe, keeping selected columns featured_col_names = X_new.columns # columns name index

如果您只需要列名：
X.columns[model.get_support()]