K-Means Clustering
An unsupervised method that groups unlabelled points into k clusters.
Syntax
KMeans(n_clusters=3)k-means is unsupervised: given data with no labels, it finds k groups. Each group has a centroid (its centre).
The algorithm
- Place
kcentroids at random. - Assign each point to its nearest centroid.
- Move each centroid to the mean of its points.
- Repeat 2-3 until stable.
Example
from sklearn.cluster import KMeans
import numpy as np
X = np.array([[1,1], [1.5,2], [8,8], [9,8], [8,9]])
model = KMeans(n_clusters=2, random_state=0, n_init=10)
model.fit(X)
print(model.labels_) # [0 0 1 1 1]
print(model.cluster_centers_) # the two centroidsWhen to use it
- A marketing team runs k-means on customer purchase vectors to discover five natural segments and craft tailored email campaigns for each.
- A computer-vision engineer uses k-means to quantise an image's colour palette to 16 colours for compression or artistic effect.
- A data scientist uses the elbow method on k-means inertia values to choose the optimal number of clusters before reporting findings.
More examples
Basic k-means clustering
Runs k-means on synthetic 4-cluster data, inspects the resulting cluster labels and the inertia (sum of squared distances to centroids).
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
X, _ = make_blobs(n_samples=300, centers=4, cluster_std=0.8, random_state=0)
kmeans = KMeans(n_clusters=4, random_state=0, n_init='auto').fit(X)
print('Labels (first 10):', kmeans.labels_[:10])
print('Inertia:', kmeans.inertia_.round(1))Elbow method to choose k
Plots inertia for k from 1 to 10 to find the elbow where adding more clusters stops providing meaningful reduction in within-cluster variance.
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
import matplotlib.pyplot as plt
X, _ = make_blobs(n_samples=500, centers=5, random_state=1)
inertias = []
for k in range(1, 11):
km = KMeans(n_clusters=k, random_state=0, n_init='auto').fit(X)
inertias.append(km.inertia_)
plt.plot(range(1, 11), inertias, marker='o')
plt.xlabel('Number of clusters k')
plt.ylabel('Inertia')
plt.title('Elbow Method')
plt.show()Predict cluster for new points
Assigns unseen data points to their nearest centroid using predict after the model has been fitted, showing how k-means generalises to new inputs.
from sklearn.cluster import KMeans
import numpy as np
X_train = np.random.randn(200, 2)
kmeans = KMeans(n_clusters=3, random_state=0, n_init='auto').fit(X_train)
new_points = np.array([[0.5, -1.2], [3.0, 2.1]])
print('Clusters:', kmeans.predict(new_points))
print('Centroids:\n', kmeans.cluster_centers_)
Discussion