如何获取sklearn中SelectKBest+chi2选中特征的原始列索引

使用sklearn库的chi2卡方特征选择方法开展特征筛选时，fit_transform()方法默认仅返回筛选转换后的特征子集，无法直接对应特征在原始矩阵中的列位置。小数据集可人工核对选中列，但面对千级甚至万级特征的大规模数据集时，人工定位列索引效率极低且容易出错。

测试示例代码 #

以鸢尾花数据集做演示，初始代码如下：

# 导入依赖库
from sklearn.datasets import load_iris
from sklearn.feature_selection import SelectKBest, chi2

# 加载数据
iris = load_iris()
X = iris.data
y = iris.target

# 将特征转换为整型分类数据
X = X.astype(int)

# 初始化卡方选择器，筛选2个和目标相关性最高的特征
chi2_selector = SelectKBest(chi2, k=2)
X_kbest = chi2_selector.fit_transform(X, y)

上述示例中人工核对可知筛选结果对应原始特征的第3、4列，但该方法无法适配大规模特征场景。

解决方案 #

拟合完成的特征选择器内置get_support()方法，可直接获取选中特征对应原始矩阵的列索引，无需人工比对。

具体用法 #

# 传入indices=True直接返回选中特征的原始列索引（索引从0开始计数）
selected_cols = chi2_selector.get_support(indices=True)
print(selected_cols)
# 鸢尾花示例输出: [2 3]，对应人工核对的第3、4列结果

# 可直接用该索引从原始特征矩阵提取筛选后的特征，结果和fit_transform输出完全一致
X_selected = X[:, selected_cols]

参数说明 #

• 不传参数调用chi2_selector.get_support()时，返回长度等于原始特征总数的布尔数组，数组元素为True代表对应位置特征被选中，为False代表未被选中
• 传入indices=True时，直接返回所有被选中特征的整数索引数组
• 该方法为sklearn特征选择类的通用接口，除SelectKBest搭配chi2的场景外，SelectPercentile、方差阈值选择等其他特征选择方法拟合后均可通过该方法获取选中列的原始索引

提示：Python序列索引默认从0开始计数，返回的索引值2、3对应日常口语表述的第3、4列，和人工核对结果一致。

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