k-means clustering

The big idea

You run a shop and have spending data for 10,000 customers, but no labels. Are there natural groups, such as bargain hunters, weekend splurgers and loyal regulars? ClusteringGrouping data points so that points in the same group are more similar to each other than to points in other groups, without using labels. finds such groups without being told what they are.

k-meansA clustering algorithm that partitions data into k groups, each represented by its mean (centroid), by alternating assignment of points to the nearest centroid and recomputation of centroids. is the simplest and most used clustering algorithm. Picture k pizza delivery depots in a new city. Each house orders from its nearest depot. Then each depot moves to the centre of the houses it serves. Some houses now have a different nearest depot, so they switch. Repeat until no house switches. The depots end up in the middle of natural neighbourhoods.

How it works

Step 1: Choose k and place k starting centroids

A centroidThe centre of a cluster: the mean of the points assigned to it. is a cluster's centre. You choose k, the number of clusters, in advance, and the algorithm needs k starting positions.

Step 2: Assign every point to its nearest centroid

Compute the (usually Euclidean) distance from each point to each centroid and give the point the label of the closest one. This carves space into regions: every location belongs to the nearest centroid.

Step 3: Move each centroid to the mean of its points

Average the coordinates of the points in each cluster. That mean is the position that minimises the cluster's sum of squared distances.

Step 4: Repeat until assignments stop changing

The quantity k-means minimises is inertiaThe sum of squared distances from each point to its assigned centroid; the objective k-means minimises. Also called within-cluster sum of squares.:

where mu_c(i) is the centroid of point i's cluster. The assign step lowers it (each point picks its closest centroid) and the update step lowers it (the mean is the best centre), so it never goes up and the loop must stop. But it stops at a local minimumA solution that cannot be improved by small changes but is worse than the best possible (global) solution., which depends on the starting centroids.

Step 5: Start well with k-means++

Random starting positions can put two centroids in one true cluster and none in another. k-means++An initialization for k-means that picks the first centroid at random from the data and each next one with probability proportional to its squared distance from the nearest centroid already chosen. spreads the starts out: pick one data point at random, then pick each next start with probability proportional to its squared distance from the nearest start so far. Far-away points are likely picks. In the code below, single runs with random starts got stuck in a worse solution 43 percent of the time; with k-means++ none did. Libraries also run several starts (n_init) and keep the lowest inertia.

Step 6: Choose k

Inertia always falls as k grows (with k = n, every point is its own cluster and inertia is 0), so you cannot just minimise it.

• Elbow method: plot inertia against k and look for the bend where adding clusters stops paying. In the code, inertia drops from 5,018 at k = 3 to 913 at k = 4, then only to 826 at k = 5: the elbow is at 4.
• Silhouette score: for each point, compare a = mean distance to its own cluster with b = mean distance to the nearest other cluster: s = (b - a) / max(a, b). It ranges from -1 to 1 and higher is better. The average silhouette peaks at 0.830 for k = 4 in the code.
• Usefulness: often k is a business decision, such as five customer segments the marketing team can act on.

Try it

Try this:

1. With k = 4 and random initialization, press Step repeatedly and watch the centroids glide. Then press New start a few times. Do you find a start where two centroids share one blob? What is the final inertia then?
2. Switch to k-means++ and repeat New start plus Run to end. How often does it end badly now?
3. Set k to 3, then 6. With 6, how does k-means split the four blobs, and how much does inertia drop?

Going deeper

Code: the loop from scratch, choosing k, initialization and a failure

import numpy as np
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs, make_moons
from sklearn.metrics import adjusted_rand_score, silhouette_score

X, truth = make_blobs(n_samples=600, centers=4, cluster_std=0.9, random_state=7)

# k-means from scratch: assign, update, repeat until nothing changes
def kmeans(X, k, seed):
    rng = np.random.default_rng(seed)
    C = X[rng.choice(len(X), k, replace=False)]
    for it in range(1, 100):
        labels = np.argmin(((X[:, None] - C[None]) ** 2).sum(-1), axis=1)
        newC = np.array([X[labels == j].mean(0) if np.any(labels == j) else C[j] for j in range(k)])
        if np.allclose(newC, C):
            break
        C = newC
    return labels, ((X - C[labels]) ** 2).sum(), it

labels, inertia, iters = kmeans(X, 4, seed=0)
print(f"from scratch: converged in {iters} iterations, inertia {inertia:.1f}, "
      f"agreement with true groups (ARI) {adjusted_rand_score(truth, labels):.3f}")

# Choosing k: inertia always falls; silhouette peaks at a good k
print("\nk   inertia   silhouette")
for k in range(2, 8):
    km = KMeans(k, n_init=10, random_state=0).fit(X)
    print(f"{k}   {km.inertia_:7.1f}   {silhouette_score(X, km.labels_):.3f}")

# Initialization: one run each from 100 different seeds
for init in ["random", "k-means++"]:
    runs = [KMeans(4, init=init, n_init=1, random_state=s).fit(X).inertia_ for s in range(100)]
    bad = np.mean(np.array(runs) > min(runs) * 1.05)
    print(f"\n{init:<9}: best inertia {min(runs):.1f}, worst {max(runs):.1f}, "
          f"runs stuck in a worse solution {bad:.0%}", end="")
print()

# A limitation: non-round clusters
Xm, ym = make_moons(400, noise=0.06, random_state=0)
print(f"\ntwo moons: k-means ARI {adjusted_rand_score(ym, KMeans(2, n_init=10, random_state=0).fit_predict(Xm)):.3f}")

from scratch: converged in 3 iterations, inertia 913.2, agreement with true groups (ARI) 1.000

k   inertia   silhouette
2   18162.5   0.594
3    5018.2   0.773
4     913.2   0.830
5     826.1   0.696
6     745.7   0.571
7     663.4   0.467

random   : best inertia 913.2, worst 13985.8, runs stuck in a worse solution 43%
k-means++: best inertia 913.2, worst 913.2, runs stuck in a worse solution 0%

two moons: k-means ARI 0.274

The from-scratch version is about ten lines and recovers the four true groups exactly (adjusted Rand index 1.000, a score of agreement between two labelings where 1 means identical). The last line shows the classic failure: on two interleaved moons k-means scores only 0.274.

Assumptions and limits

• Round, similar-sized clusters. Using distance to the mean implies clusters are roughly spherical and of similar spread. Elongated, nested or moon-shaped clusters are cut in the wrong places; use DBSCAN or Gaussian mixtures ( Hierarchical clustering and DBSCAN , Gaussian mixture models and EM ).
• Scale matters. A feature measured in dollars swamps one measured in years. Standardize features first ( Feature engineering ).
• Every point is assigned. Outliers get a cluster anyway and drag its centroid. k-medians or k-medoids are more robust.
• Cost is about n x k x d per iteration, so k-means scales to millions of points; mini-batch k-means updates centroids from small random batches for even larger data.

Trade-offs

Common mistakes

• Running k-means on unscaled features, so one large-valued feature decides everything.
• Using a single random initialization and trusting the result.
• Picking k by minimising inertia, which always favours more clusters.
• Applying k-means to clearly non-spherical clusters and trusting the boundaries.
• Interpreting clusters as real-world categories without checking them against domain knowledge.

In the wild

k-means turns up in customer segmentation, grouping documents or products by their embeddings, compressing images by reducing colours to k palette entries, and as a building block elsewhere: vector databases often use it to partition embeddings into buckets for fast approximate search (the inverted-file index behind many retrieval systems), and speech and vision pipelines use it to build codebooks of typical feature vectors.

Check yourself

Recap