K-means聚类后提取数据构建决策树的高效实现方案问询

嘿，我来帮你把这个繁琐的流程简化一下！你现在遇到的X与Y值不匹配问题，核心是手动遍历索引拼接数据容易出错，而利用pandas的分组功能，结合K-means标签与训练数据的对应关系，就能完美解决这个问题，同时让代码简洁很多。

核心思路 #

kmeans.labels_的顺序和X_train的样本顺序是完全对应的——第一个标签对应X_train的第一行样本，以此类推。所以我们可以把聚类标签、特征数据X_train和类别数据y_train合并成一个DataFrame，然后按聚类标签分组，就能直接获取每个簇对应的特征和类别，完全不会出现匹配错误。

优化后的完整代码 #

from sklearn.cluster import KMeans
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
import pandas as pd

# 你的原始数据准备部分（保留你的逻辑）
X = clusterDF[clusterDF.columns[clusterDF.columns.str.contains("'AB'")]]
y = clusterDF['Class']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.30, random_state=1)

# K-means聚类部分（保留你的参数，注意K-means是无监督算法，fit不需要传入y_train）
kmeans = KMeans(n_clusters=3, n_init=5, max_iter=3000, random_state=1)
kmeans.fit(X_train)

# 关键：将聚类标签与训练数据合并，保证样本一一对应
train_clustered = pd.concat(
    [
        X_train,
        y_train,
        pd.Series(kmeans.labels_, name='cluster_label', index=X_train.index)
    ],
    axis=1
)

# 遍历每个簇，构建决策树
for cluster_id, cluster_data in train_clustered.groupby('cluster_label'):
    print(f"\n\t-----Classifier {cluster_id + 1}----")
    
    # 提取当前簇的特征和类别
    X_cluster = cluster_data.drop(['cluster_label', 'Class'], axis=1)
    y_cluster = cluster_data['Class']
    
    # 划分簇内的训练/测试集（注意：如果簇样本量过小，可调整test_size或跳过）
    Xc_train, Xc_test, yc_train, yc_test = train_test_split(
        X_cluster, y_cluster, test_size=0.30, random_state=1
    )
    
    # 训练决策树并获取结果（复用你的getResults函数）
    classifier = DecisionTreeClassifier(criterion='entropy', max_depth=2)
    classifier = getResults(
        Xc_train, yc_train, Xc_test, yc_test,
        classifier, filename=f'classif_{cluster_id + 1}'
    )

为什么这个方法更好？ #

• 完全避免匹配错误：所有数据来自同一个DataFrame，样本的特征、类别和聚类标签天然一一对应，不需要手动索引校验
• 代码极度简化：用groupby替代了复杂的循环拼接逻辑，行数减少一半以上，可读性和维护性大幅提升
• 灵活可控：可以轻松查看每个簇的样本量（比如cluster_data.shape[0]），如果簇太小可以直接跳过，避免后续训练出现问题

额外注意事项 #

• 如果某个簇的样本数量很少（比如小于10），test_size=0.3可能会导致测试集样本不足，建议根据实际情况调整test_size，或者在分组后先判断len(cluster_data)再决定是否训练
• 确保X_train和y_train的索引一致——train_test_split默认会保持索引对应，如果你之前修改过索引，可以用reset_index(drop=True)统一两者的索引

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