是否可以根据kmeans结果的注释而不是聚类结果对其进行着色？

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans

# 假设X是您的数据集，annotations是对应的注释
X = np.array([[1, 2], [1, 4], [1, 0],
              [10, 2], [10, 4], [10, 0]])
annotations = np.array(['A', 'A', 'A', 'B', 'B', 'B'])  # 示例注释

# 执行K-means聚类
kmeans = KMeans(n_clusters=2)
clusters = kmeans.fit_predict(X)

# 根据注释着色
unique_annotations = np.unique(annotations)
colors = plt.cm.jet(np.linspace(0, 1, len(unique_annotations)))

plt.figure(figsize=(8, 6))
for annotation, color in zip(unique_annotations, colors):
    subset = X[annotations == annotation]
    plt.scatter(subset[:, 0], subset[:, 1], color=color, label=f'Annotation {annotation}')

# 可以选择性地添加聚类中心
centers = kmeans.cluster_centers_
plt.scatter(centers[:, 0], centers[:, 1], c='black', s=200, alpha=0.5, label='Centroids')

plt.legend()
plt.show()