k-Means Clustering

A product manager hands you six customers with two features each — visits per month and average basket size — and asks for two segments the growth team can target differently. k-Means answers with Lloyd's algorithm: a two-step loop simple enough to run on a whiteboard and effective enough to remain the clustering baseline half a century after publication. The two steps are the entire algorithm, and being able to trace them is the difference between using k-Means and debugging it.

Start with two points chosen as initial centroids, mu_1 = (1, 2) and mu_2 = (5, 8). The assignment step gives every point to its nearest centroid under Euclidean distance. Squared distances suffice — the square root is monotone, so the argmin is unchanged — and skipping the root saves real work:

The update step recomputes each centroid as the mean of its assigned points. The mean is not a convention: among every candidate center, the arithmetic average is the unique point that minimizes the sum of squared distances to the cluster's members — take the gradient of that sum with respect to the center, set it to zero, and the mean falls out. Any other choice (median, midpoint) would leave the objective improvable, and the stopping argument below would collapse. Now mu_1 = (1.25, 1.9) and mu_2 = (7.875, 9).

Reassign with the new centroids. Only p3 (5, 8) sits near a boundary: 51.3 squared from mu_1, 9.27 from mu_2, so it stays in C2. Every other label is unchanged, and the loop stops exactly when a full assignment round produces no label changes. Two iterations, six points.

The quantity being driven down is inertia, the within-cluster sum of squares:

inertia = sum over clusters C of ( sum over x in C of ||x - mean(C)||^2 )

At convergence this instance sits at roughly 0.29 + 19.2 = 19.5. sklearn exposes it as KMeans.inertia_, and the next topic leans on it heavily — but notice what it actually measures: compactness in feature space, nothing about whether the two segments are useful to a human.

A faithful implementation in under thirty lines:

import numpy as np

def lloyd(X, k, iters=100):
    rng = np.random.default_rng(0)
    C = X[rng.choice(len(X), k, replace=False)]   # random initial centroids
    labels = None
    for _ in range(iters):
        new_labels = ((X[:, None, :] - C[None]) ** 2).sum(-1).argmin(1)
        C = np.array([X[new_labels == j].mean(0) for j in range(k)])
        if np.array_equal(new_labels, labels):
            return new_labels, C                  # a full round changed nothing
        labels = new_labels
    return labels, C

Being asked to "write k-Means" in an interview usually means exactly this loop, so practice one hand iteration on a small example before the screen. The mechanics here are also the raw material for Choosing the Cluster Count, which is mostly about reading the inertia this loop produces.

k-Means carries a termination guarantee that surprises people who expect iterative algorithms to oscillate — followed by an optimality surprise in the other direction. It always stops. It can stop at a bad answer.

The guarantee comes from tracking the objective through both steps. The assignment step moves each point to its nearest centroid, so the sum of squared distances cannot increase; the update step then places every center at the best possible location for its assigned points, so the sum drops again unless the centroids were already optimal. Inertia is monotonically non-increasing. Only finitely many partitions of n points into k groups exist, and the objective can never revisit a higher value, so the loop cannot cycle: it must reach a labeling that no reassignment improves. That is why a from-scratch implementation can stop on "labels unchanged", while sklearn uses two cheaper proxies: a relative inertia change below tol (default 1e-4) or max_iter (default 300) exhausted.

The optimality gap matters just as much. Finding the best partition is NP-hard, so every run settles into a local minimum of inertia, and which one depends on where the centroids started:

inertias = []
for seed in range(5):
    km = KMeans(n_clusters=4, init='random', n_init=1, random_state=seed)
    inertias.append(round(km.fit(X).inertia_))
# e.g. [4821.7, 4790.3, 4318.9, 4825.1, 4299.4]  — three distinct answers

A 12% spread in inertia across seeds on a toy dataset is unremarkable, and it means a single k-Means run is not a result you can report to anyone.

k-means++ is the standard fix, and it works by making the initial centroids far apart on purpose. The first centroid is drawn uniformly at random; each subsequent one is drawn with probability proportional to D(x)^2, where D(x) is the distance from x to the nearest already-chosen centroid. Points far from existing centroids become likely seeds; crowded regions do not. Seeding costs one extra O(n·k·d) pass, and the theory gives a competitive-ratio bound of O(log k) against the optimum — not optimal, but no longer at the mercy of one bad first draw.

In production you combine the ideas and stop thinking about them:

from sklearn.cluster import KMeans

km = KMeans(n_clusters=4, init='k-means++', n_init=10, random_state=42)
labels = km.fit_predict(X_scaled)
km.inertia_   # inertia of the best of the 10 restarts
km.n_iter_    # iterations the winning run needed

Read those defaults as a confession: sklearn runs ten independent k-means++ starts and keeps the lowest inertia because one start is not trustworthy. random_state makes the pick reproducible; leave it unset and your labels shift run to run, which makes every downstream metric wobble.

Two operational details round it out. Per-iteration cost is O(n·k·d), so k-Means scales linearly in the number of points — rare among ML algorithms — and MiniBatchKMeans trades a small inertia penalty (often well under 1%) for large speedups past a few hundred thousand rows. The algorithm parameter ('lloyd' by default, or 'elkan', which uses triangle inequalities to skip distance computations) changes runtime, never the objective.

Inertia cannot increase after a correctly completed Lloyd update for a fixed initialization: assigning each point to its nearest center and replacing each center with its cluster mean each minimizes the respective part of the squared-distance objective. It can still stop at a poor local minimum because this objective is non-convex in the assignments and centers jointly. k-means++ spreads starting centers using distance-weighted sampling and usually improves the starting objective, but it does not make the solution globally optimal. Run multiple initializations, compare inertia, and inspect whether the resulting assignments are substantively stable. A stable low objective is useful evidence; it is not proof that the clusters match real categories.

Because k-Means always returns k clusters plus a tidy inertia number, a structurally wrong answer looks exactly like a good one. The algorithm has no notion of failure. Your job is to recognize the four conditions where its output is wrong by construction.

Unscaled features break the distance metric. k-Means optimizes squared Euclidean distance, so the largest-magnitude feature owns the objective. A customer table with income in dollars (~10^4) and visits per month (~10^1) squares that gap: income differences contribute on the order of 10^8 while visits contribute ~10^2 — a factor of a million. The "clusters" become income brackets and visits become noise. The fix is unglamorous and non-negotiable: fit a StandardScaler (or RobustScaler when outliers are known) on the features, cluster the scaled matrix, and keep that same scaler for every downstream step.

Spherical clusters are assumed, never checked. The objective only measures distance to a center, so it answers "where are the centers" even when the truth is two interleaving crescents. On that classic shape k-Means slices each crescent in half, and nothing in the inertia curve objects — inertia improves with every wrong cut. A scatter plot of your two most informative features showing long, curved, or wildly unequal-size groups is the visual evidence to switch algorithms, not to raise k.

Outliers move centroids. A single point at (500, 500) in a table that otherwise lives near the origin can drag a whole centroid toward itself and sacrifice every genuine member of that cluster. Because the center is a mean, k-Means has no breakdown point: a handful of extremes suffices. Remove or clip known outliers, or use k-medoids, whose centers are real observations.

Categorical inputs have no Euclidean distance. One-hot encoding every category technically works, but it inflates dimensionality and reduces "distance" to a count of mismatches; k-prototypes (k-Means with a Hamming-style cost for categorical columns) is the standard answer for mixed tables.

One tie back to The Kernel Trick in SVMs: the remedy of clustering in a feature space where groups become rounder is the same idea as the RBF kernel — Euclidean distance only expresses what the geometry you gave it can express.

The decision rule to carry into an interview: k-Means is the right default when features are numeric and scaled, groups plausibly look roughly round, and speed at scale matters. When those hold, nothing beats it. When they do not, hierarchical clustering and DBSCAN exist for exactly those cases.

Before interpreting a cluster, compare its profile with the overall population and with plausible baseline segments. A group defined mostly by one extreme numeric feature may be an artifact of scale or an outlier rather than a useful segment. Check assignment stability after resampling rows or changing the seed, and test whether the group persists after reasonable preprocessing choices. Do not name clusters from their centroid alone: inspect distributions because the mean vector can describe no actual member when features trade off. For product decisions, validate the segment on a later time window and ask whether a distinct action for that group improves an outcome.