2017-04-08 115 views
0

我想要选择基于树的功能。 我的数据集有大约30列,做完之后,大约有5个。 对我来说,这对我来说很好,我遇到的问题是我得到的5列的数据集不保留列的名称,我可以不识别他们。Scikit学习基于树的功能选择保留列名称?

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) 

有人可以帮我吗?

+0

试试['get_support '](http://scikit-learn.org/stable/modules/generated/sklearn.feature_selection.SelectFromModel.html#sklearn.feature_selection.SelectFromModel.get_support)功能。正如'model.get_support(indices = True)' – ehudk

+0

你好ehudk,按照你所告诉我的,我得到所选原始数据集的列号。确实没有表明列的名称,但它是一个解决方案。如果有人愿意贡献,我会留下这个问题。非常感谢您的帮助 – Javi

回答

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()]

+0

感谢ehudk我会尽力的! – Javi