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

# set up some constants
num_clusters = 4
np.random.seed(0)
X = np.random.rand(100, 2)

# we have to set n_init explicitly now or we get nasty warning.
kmeans = KMeans(n_clusters=num_clusters, n_init=10, random_state=0)
kmeans.fit(X)

# Get the cluster labels and centroids from the semiprivate fields
centroids = kmeans.cluster_centers_
labels = kmeans.labels_

# plot the position of the centroids and the data points around it
plt.figure(figsize=(8, 6))
plt.scatter(X[:, 0], X[:, 1], c=labels, cmap='viridis')
plt.scatter(centroids[:, 0], centroids[:, 1], marker='x', color='red', s=200)
plt.title('k-means cluster with 4 centroids')
plt.xlabel('X')
plt.ylabel('Y')
plt.show()