仅使用感兴趣的列进行聚类(如您的示例中所示)。然后添加标签列表kmeans.labels_
作为另一列X_train
(or X_test
)。标签的顺序与原始行的顺序相同。
# A toy DF
X = pd.DataFrame({'id': [1,2,3,4,5],
'value_1': [1,3,1,4,5],
'value_2': [0,0,1,5,0]})
# Split ALL columns
(X_train, X_test) = train_test_split(X,test_size=0.30)
# Cluster using SOME columns
kmeans = kmeans.fit(X_train[['value_1','value_2']])
# Save the labels
X_train.loc[:,'labels'] = kmeans.labels_
既然两者X_train
and X_tests
是切片X
,您可能会在此处看到警告:
尝试在 DataFrame 的切片副本上设置一个值。
你可以忽略它。
X_train
# id value_1 value_2 labels
#4 5 5 0 0
#0 1 1 0 0
#3 4 4 5 1