K-means clustering

1. Intuition

You run a delivery company and want to open three depots so that, on average, customers are close to one. You don't know where the depots should go, so you guess three spots. Each customer is served by the nearest depot. Then you move each depot to the center of the customers it serves. Some customers now have a different nearest depot, so you reassign them, and move the depots again. After a few rounds, nothing changes: you have found three clusters and their centers.

Two things are worth noticing:

• You picked K. The algorithm never decides how many groups exist. It finds the best K groups it can, even if the data has no real groups at all.
• Your first guess matters. Start with two depots in one city and one in the far countryside, and the process may settle into a poor arrangement it cannot escape. This is a local optimum.

2. Visualization

Try it (predict first, then check):

1. Load Rings and predict whether any choice of 2 centroids can separate the inner ring from the outer one. Run it.
2. On Blobs, change the seed a few times with K = 3. Do all runs end at the same inertia? Why not?

What to notice:

• Inertia (the total squared distance to centers) always decreases as K grows, so "lowest inertia" would pick K = n. Look for the elbow where adding clusters stops helping much.
• Silhouette rewards clusters that are tight and far apart, so it can peak and then fall.
• K-means splits the plane into regions around each center (a Voronoi diagram). Any cluster that is not roughly round and convex will be cut wrongly.

3. The math

Symbols

Objective

$\lVert a \rVert^2$ is the squared Euclidean length: $\sum_j a_j^2$.

The two steps (Lloyd's algorithm)

Assignment: with centroids fixed, give each point to its nearest centroid.

Update: with assignments fixed, move each centroid to the mean of its points.

Why does $J$ never increase? The assignment step picks, for each point separately, the centroid that minimizes its term, so no term can grow. The update step uses the fact that for a fixed set of points, the single location minimizing the sum of squared distances is their mean (set the derivative $\sum_{i \in C_k} -2(x_i - \mu_k)$ to zero and solve). Since $J$ never increases and there are finitely many possible assignments, the algorithm always stops, though possibly at a local optimum.

Worked example

Six points, $K = 2$. Initialize the centroids at A and F: $\mu_1 = (1, 1)$, $\mu_2 = (4.5, 5)$.

Step 1: squared distances to each centroid. For C to $\mu_1$: $(3-1)^2 + (4-1)^2 = 4 + 9 = 13$. To $\mu_2$: $(3-4.5)^2 + (4-5)^2 = 2.25 + 1 = 3.25$.

Step 2: inertia after assignment. Add each point's distance to its own centroid: $0 + 1.25 + 3.25 + 4.25 + 1.00 + 0 = 9.75$.

Step 3: update.

Step 4: inertia after update. For cluster 1: A $(0.25^2 + 0.5^2 = 0.3125)$, B $(0.25^2 + 0.5^2 = 0.3125)$. For cluster 2: C $(1 + 1.5625 = 2.5625)$, D $(1 + 3.0625 = 4.0625)$, E $(0.25 + 0.0625 = 0.3125)$, F $(0.25 + 0.0625 = 0.3125)$. Total $= 7.875$, down from 9.75.

Step 5: assign again. Every point keeps its cluster (C is still much closer to $\mu_2$: 2.56 versus 9.31), so the centroids no longer move. K-means has converged in one update.

4. Implementation

Build it from scratch

Four small functions in NumPy cover the whole algorithm. The full script is code/07-unsupervised/kmeans.py.

Assignment step. Broadcasting X[:, None, :] - mu[None, :, :] makes an [n, K, p] array of differences, so one line gives every point's squared distance to every centroid. argmin along the centroid axis picks the nearest one.

def assign(X, mu):
    """Assignment step: squared distance to every centroid, then the nearest one."""
    d2 = ((X[:, None, :] - mu[None, :, :]) ** 2).sum(axis=2)   # [n, K] squared distances
    return d2, d2.argmin(axis=1)                               # [n, K], [n]

Predict: point C in the worked example is at squared distance 13 from $\mu_1$ and 3.25 from $\mu_2$. What does argmin return for C's row, given that index 0 means cluster 1?

Objective and update step. mu[a] looks up each point's own centroid, so the inertia is one subtraction and a sum. The update replaces each centroid with the mean of its points. A centroid that loses all its points would produce NaN from an empty mean, so it stays where it is.

def inertia(X, mu, a):
    """Objective J: total squared distance from each point to its own centroid."""
    return float(((X - mu[a]) ** 2).sum())

def update(X, a, mu):
    """Update step: move each centroid to the mean of its points (keep it if the cluster is empty)."""
    return np.array([X[a == j].mean(axis=0) if np.any(a == j) else mu[j] for j in range(len(mu))])

Predict: after an update, can inertia(X, new_mu, a) ever be larger than inertia(X, mu, a) with the same assignments a?

The loop (Lloyd's algorithm). Start from K random data points, then alternate the two steps until the centroids stop moving. This version uses one random start, which is exactly the weakness k-means++ and n_init fix.

def kmeans(X, k, seed=0, iters=100):
    """Lloyd's algorithm from a random start: alternate assign and update until centroids stop moving."""
    rng = np.random.default_rng(seed)
    mu = X[rng.choice(len(X), k, replace=False)]               # [k, p] k random data points
    for _ in range(iters):
        _, a = assign(X, mu)
        new = update(X, a, mu)
        if np.allclose(new, mu):
            break
        mu = new
    return mu, a

Predict: two of the four random starting points land in the same true blob. Will the loop ever move one of them to the empty blob?

Silhouette score. For each point, $a$ is its mean distance to the rest of its own cluster and $b$ is its mean distance to the nearest other cluster; $s = (b - a)/\max(a, b)$. The [n, n] distance matrix makes this $O(n^2)$ memory, fine for 300 points but not for a million.

def silhouette(X, a):
    """Mean silhouette: s = (b - a) / max(a, b) per point, a = own-cluster mean distance, b = nearest other cluster."""
    D = np.sqrt(((X[:, None, :] - X[None, :, :]) ** 2).sum(axis=2))   # [n, n] pairwise distances
    s = np.zeros(len(X))
    for i in range(len(X)):
        own = a == a[i]
        if own.sum() == 1:                                     # singleton cluster: s = 0 by convention
            continue
        a_i = D[i, own].sum() / (own.sum() - 1)                # exclude the zero distance to itself
        b_i = min(D[i, a == c].mean() for c in np.unique(a) if c != a[i])
        s[i] = (b_i - a_i) / max(a_i, b_i)
    return float(s.mean())

Predict: a point sits exactly halfway between its own centroid and another cluster, so $a \approx b$. What is its silhouette, and what does a negative value mean?

Use the library

scikit-learn's KMeans defaults to k-means++ initialization and runs several starts (n_init), keeping the lowest inertia. It computes distances with matrix products and runs in compiled code, so it is much faster than the loop above.

from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
from sklearn.preprocessing import StandardScaler

X_scaled = StandardScaler().fit_transform(X)
km = KMeans(n_clusters=4, init="k-means++", n_init=10, random_state=0).fit(X_scaled)
print(km.inertia_, silhouette_score(X_scaled, km.labels_))

Run it

From guide/:

uv run code/07-unsupervised/kmeans.py

Key lines of output:

Iteration 1: centroids [[1.0, 1.0], [4.5, 5.0]]
  inertia after assignment: 9.750
  update -> [[1.25, 1.5], [4.0, 5.25]]  inertia 7.875
  ...
  converged

from scratch, one random start (seed 1): inertia=2359.7  <- stuck in a local optimum
from scratch, best of 10 random starts:  inertia=462.8
scikit-learn (k-means++, n_init=10):      inertia=462.8
silhouette at K = 4: from scratch 0.8265, scikit-learn 0.8265
K, inertia, silhouette: [(1, 22193.5, None), (2, 8941.6, 0.599), (3, 2402.8, 0.776), (4, 462.8, 0.826), (5, 414.0, 0.721), ...]

The first block reproduces the worked example. On the synthetic blobs, a single random start got stuck at inertia 2359.7 (two true blobs merged, one split). The best of 10 random starts and scikit-learn both reach 462.8. That is why n_init exists. The from-scratch silhouette matches scikit-learn, and both the elbow and the silhouette peak point to K = 4.

5. Engineering

Initialization. k-means++ picks the first centroid at random, then picks each next centroid with probability proportional to its squared distance from the nearest centroid already chosen. That spreads the starts out and usually avoids bad local optima. Still run several initializations (n_init) and keep the lowest inertia.

Choosing K. Combine signals: the elbow in inertia, the silhouette score (from −1 to 1; higher means tighter, better-separated clusters), stability across seeds and data subsamples, and above all whether the clusters are useful and interpretable to the people who will act on them.

Preprocessing. Scale features, because K-means uses Euclidean distance (the same issue as KNN). Consider PCA first for very high-dimensional data. One-hot categorical features distort means; use k-modes or a different method for mostly categorical data.

Cost. Each iteration costs $O(nKp)$. It is fast and scales to millions of points; MiniBatchKMeans updates centroids from small random batches for even larger data.

Where it breaks. K-means assumes roughly spherical, similarly sized, similarly dense clusters. It fails on elongated or curved shapes, very different cluster sizes, and outliers (one far point drags a centroid). Use DBSCAN for arbitrary shapes and noise, Gaussian mixtures for elliptical clusters and soft membership, and hierarchical clustering when you want nested structure.

Key distinction

A cluster is not a category. K-means always returns K groups, even for uniform random data. A discovered cluster is a hypothesis about structure. Validate it with domain experts, stability checks, and downstream usefulness before naming it "high-value customers."

Common mistakes

• Forgetting to scale features, so one large-unit feature decides every cluster.
• Choosing K by the lowest inertia (which always favors more clusters).
• Running a single initialization and trusting the result.
• Treating cluster labels as stable IDs across retrains; the numbering is arbitrary and changes between runs.
• Applying K-means to data whose groups are not convex and concluding "there is no structure."

6. Knowledge check

Practice exercise. Points on a line: 1, 2, 3, 10, 11, 12. Start with $\mu_1 = 1$, $\mu_2 = 2$. Run K-means until it converges. What are the final centroids and inertia?

Implementation challenge. Implement k-means++ initialization for the from-scratch version, then run both initializations 50 times on the synthetic blobs and compare the distribution of final inertias.

Summary

• K-means minimizes the total squared distance from points to their assigned centroids.
• It alternates two steps that can only lower the objective: assign to the nearest centroid, move centroids to means.
• It converges to a local optimum; use k-means++ and multiple initializations.
• Choose K with the elbow, silhouette, stability, and domain judgment, never inertia alone.
• It assumes round, similar-size clusters and scaled features; a discovered cluster is a hypothesis, not a fact.

Next: DBSCAN, hierarchical clustering, and Gaussian mixtures

Related: K-nearest neighbors · PCA · Model card: K-means

Parts of this lesson are adapted from AI Engineering from Scratch by Rohit Ghumare (MIT license, notice): Unsupervised Learning.

Interview angle