Extracting centroids using k-means clustering in python?

Question

I have some data in a 1D array with shape [1000,] with 1000 elements in it. I applied k-means clustering on this data with 10 as number of clusters. After applying the k-means, I got cluster labels (id's) with shape [1000,] and centroids of shape [10,] for each cluster. The labels array allots value between 0 and 9 to each of the 1000 elements. However, I want each element to show its centroid rather than its cluster id's. How can I achieve this?

from sklearn.cluster import KMeans
kmeans = KMeans(n_clusters=10)
kmeans.fit(data)   #data is of shape [1000,]
#learn the labels and the means
labels = kmeans.predict(data)  #labels of shape [1000,] with values 0<= i <= 9
centroids  = kmeans.cluster_centers_  #means of shape [10,] 

In the code above, I want respective centroids for each element in the [1000,] array instead of its cluster id's.

Sheenasheeny

Answer 1

Seems like a list comprehension would work well.

centroid_labels = [centroids[i] for i in labels]

Extortioner

Answer 2

Just use the centroids array as lookup table:

samplesCentroids = centroids[labels]

Faker