kmeans.py

from matplotlib import pyplot as plt
from sklearn.datasets import make_blobs
from sklearn.cluster import KMeans

X, y = make_blobs(n_samples=500, centers=20, random_state=999)
plt.scatter(X[:,0], X[:,1])

kmeans = KMeans(n_clusters=4, init='k-means++', max_iter=300, n_init=10)
y_preds = kmeans.fit_predict(X)
centroids = kmeans.cluster_centers_
plt.scatter(X[:, 0], X[:, 1], c=y_preds)
plt.scatter(centroids[:, 0], centroids[:, 1], s=50, c='red')