如何在scikit学习列选择器管道中只选择几个列？

from sklearn.datasets import load_iris
from sklearn.feature_selection import SelectKBest, f_classif

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

# 假设我们只想选择第1列和第3列
columns_to_select = [0, 2]

# 创建SelectKBest对象，但只选择我们感兴趣的列
selector = SelectKBest(score_func=f_classif, k=len(columns_to_select))
X_new = selector.fit_transform(X[:, columns_to_select], y)

print(X_new)

from sklearn.datasets import load_iris
from sklearn.compose import ColumnTransformer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import FunctionTransformer

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

# 假设我们只想选择第1列和第3列
columns_to_select = [0, 2]

# 创建一个函数来选择特定的列
def select_columns(X, cols):
    return X[:, cols]

# 创建ColumnTransformer对象
column_transformer = ColumnTransformer(
    transformers=[
        ('select', FunctionTransformer(select_columns, kw_args={'cols': columns_to_select}), columns_to_select)
    ])

# 创建一个包含ColumnTransformer的管道
pipeline = Pipeline(steps=[('select_cols', column_transformer)])

# 应用管道
X_new = pipeline.fit_transform(X)

print(X_new)