k-means clustering

k-means finds groups in data that has no labels. You tell it how many groups to look for, and it places one centre per group so that every point is as close as possible to its nearest centre. It is the first clustering algorithm to learn and, for many problems, the only one you need.

The algorithm

1. Choose $k$ and place $k$ centres at random (or, better, spread out using the k-means++ rule).
2. Assign: give each point to its nearest centre.
3. Update: move each centre to the mean of the points assigned to it.
4. Repeat 2 and 3 until the assignments stop changing.

Each iteration lowers (or keeps) the total squared distance from points to their centres, so the algorithm always converges. It may converge to a local optimum, which is why scikit-learn runs it several times from different starts (n_init) and keeps the best.

The objective

where $\mu_{c(i)}$ is the centre assigned to point $i$. This quantity is called inertia or within cluster sum of squares.

From scratch

import numpy as np

rng = np.random.default_rng(0)
X = np.vstack([rng.normal([0, 0], 0.6, (30, 2)), rng.normal([5, 5], 0.6, (30, 2)), rng.normal([0, 5], 0.6, (30, 2))])

k = 3
centres = X[rng.choice(len(X), k, replace=False)]
for iteration in range(10):
    distances = np.linalg.norm(X[:, None, :] - centres[None, :, :], axis=2)
    labels = distances.argmin(axis=1)
    new_centres = np.array([X[labels == j].mean(axis=0) for j in range(k)])
    if np.allclose(new_centres, centres):
        print("converged after", iteration, "iterations")
        break
    centres = new_centres

print(centres.round(2))
print(np.bincount(labels))

With scikit-learn

import numpy as np
from sklearn.cluster import KMeans

rng = np.random.default_rng(0)
X = np.vstack([rng.normal([0, 0], 0.6, (30, 2)), rng.normal([5, 5], 0.6, (30, 2)), rng.normal([0, 5], 0.6, (30, 2))])

km = KMeans(n_clusters=3, n_init=10, random_state=0).fit(X)
print(km.cluster_centers_.round(2))
print("inertia:", round(km.inertia_, 2))
print("new point belongs to cluster", km.predict([[4.5, 4.8]])[0])

Choosing k: the elbow method

Inertia always falls as $k$ grows (with $k$ equal to the number of points it is zero). Plot inertia against $k$ and look for the elbow where the drop flattens out.

import numpy as np
from sklearn.cluster import KMeans

rng = np.random.default_rng(0)
X = np.vstack([rng.normal([0, 0], 0.6, (30, 2)), rng.normal([5, 5], 0.6, (30, 2)), rng.normal([0, 5], 0.6, (30, 2))])

for k in range(1, 7):
    inertia = KMeans(n_clusters=k, n_init=10, random_state=0).fit(X).inertia_
    print(f"k={k}: inertia={inertia:7.2f} " + "#" * int(inertia / 10))

The bend at $k = 3$ is clear here. Real data is rarely this tidy.

Choosing k: the silhouette score

For each point, compare its average distance to its own cluster ($a$) with its average distance to the nearest other cluster ($b$):

Values near 1 mean well separated clusters; near 0 means overlapping; negative means probably in the wrong cluster. Average over all points and pick the $k$ with the highest score.

import numpy as np
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score

rng = np.random.default_rng(0)
X = np.vstack([rng.normal([0, 0], 0.6, (30, 2)), rng.normal([5, 5], 0.6, (30, 2)), rng.normal([0, 5], 0.6, (30, 2))])

for k in range(2, 7):
    labels = KMeans(n_clusters=k, n_init=10, random_state=0).fit_predict(X)
    print(f"k={k}: silhouette={silhouette_score(X, labels):.3f}")

Practical notes

• Scale features. Distance based, like k-NN.
• Clusters are spherical. k-means assumes round groups of similar size. For elongated or nested shapes use DBSCAN or Gaussian mixtures.
• Sensitive to outliers. A far away point drags a centre towards it.
• Categorical data does not have a meaningful mean. Use k-modes or one-hot encode with care.
• Interpreting clusters is on you. Look at the centre of each cluster in original units and describe it in words: "high spend, low frequency".

Where it is used

Customer segmentation, grouping documents or images, compressing colours in an image (each pixel replaced by its cluster centre), and as a preprocessing step to create a "cluster id" feature for a supervised model.

Practice

1. Generate data with four clear groups and confirm that both the elbow and the silhouette method point to $k = 4$.
2. Add three far away outlier points and see how the centres move.
3. Cluster the iris features (ignore the labels) with $k = 3$ and compare the cluster assignments with the true species using a confusion matrix.