K-means clustering

What this chapter covers

Every model so far came with an answer key. The data had a label column, we trained to match it, and we scored ourselves on how often we did. This chapter takes the answer key away. Here are 150 points on a plane and nothing else — no classes, no truth, just coordinates. The job is to find the groups that are already there, without ever being told what a group is.

That's clustering, and k-means is the one everyone reaches for first. The idea fits in a sentence: guess where the cluster centers are, assign every point to its nearest center, move each center to the middle of the points that chose it, and repeat until nothing moves. We build it by hand in NumPy, one function at a time, then let scikit-learn build the same thing and watch the two land on the identical answer. And because the whole algorithm is just points and centers sliding around a plane, we get to watch it happen — the centroids start scattered on random points and walk, frame by frame, into the heart of each blob.

The data is four well-separated blobs from make_blobs. We do have the true blob ids, but only to grade the result at the very end. k-means never sees them. That's the line this chapter is really about: the first time in the course the model learns structure with no one holding the answer.

A bit of history

The recipe is older than the name. In 1956 the Polish mathematician Hugo Steinhaus wrote down the problem of splitting a body into parts that minimize the distance of each point to its part's center — the k-means objective, in French, with no computer to run it on. A year later, in 1957, Stuart Lloyd at Bell Labs described the exact assign-and-average iteration we use today, as a way to quantize signals for pulse-code modulation. His write-up circulated internally and wasn't formally published until 1982, which is why you'll see the algorithm called "Lloyd's algorithm" and dated to two different decades depending on who's citing it.

Edward Forgy proposed essentially the same method in 1965, so the raw scheme sometimes carries his name too. The label that stuck came from James MacQueen, who in 1967 coined the term "k-means" and framed it as the procedure we teach now. For thirty years it ran on whatever seeding you gave it, and a bad seeding could wreck it — until 2007, when David Arthur and Sergei Vassilvitskii added k-means++, a smarter way to place the initial centers that made the whole thing robust enough to be a sane default. That's the version scikit-learn ships, and the version we build toward here.

The intuition

Look at the data with no colors on it. Your eye finds the groups instantly — there's a clump up top, one on the right, one down low, one on the left. You didn't compute anything; you just saw density. k-means is a way to make a computer do that same thing, and the trick is that it never sees the groups either. It only ever knows two moves.

Guess k center points. Then loop: color each point by whichever center is closest, and slide each center to the average position of the points that are now its color. That's the entire algorithm. The centers chase the density. Points near a center pull it toward them; when a center drifts into a clump, more points claim it, which pulls it in tighter, until it settles at the balance point and stops. Do that for all k centers at once and they spread out to cover the blobs, each one parking itself in the middle of a group.

The one thing the algorithm can't do is decide how many groups there are. You hand it k. Get k right and it finds clean clusters; get it wrong and it will still hand you exactly that many, splitting a real group in half or fusing two together without complaint. Here's the raw data it's going to work on — 150 points, no labels, gray:

Four blobs, and you found them without being told there were four. Now let's make the algorithm find the same thing.

The math

Let $x_1, \dots, x_N$ be the points and $\mu_1, \dots, \mu_k$ the $k$ cluster centers, or centroids. Assign each point to its nearest centroid; write $c_i$ for the index of the centroid point $i$ belongs to. The quantity k-means minimizes is the within-cluster sum of squares, called the inertia:

Every point contributes the squared distance to its own centroid. A tight cluster contributes little; a loose one, a lot. Lloyd's algorithm drives $J$ down by alternating two steps, each of which is the exact minimizer of $J$ if you hold the other half fixed.

The assignment step fixes the centroids and picks the best cluster for every point — the nearest one:

The update step fixes the assignments and picks the best centroid for every cluster. Hold a set of points and ask which single location minimizes the sum of squared distances to them; the answer is their mean, which is why the step is just an average:

Both steps can only lower $J$ or leave it unchanged, and there are finitely many ways to assign $N$ points to $k$ clusters, so the loop has to stop. That's the whole guarantee — and also the whole catch. It converges, but to a local minimum of $J$, and which one depends entirely on where the centroids started. That single fact is why seeding gets its own function later.

What it's good at, what it isn't

The good side is why it's everywhere. It's fast — each iteration is one distance matrix and one group-by-average, and it usually converges in a handful of passes, so it scales to data that makes fancier methods choke. It's simple enough to implement in an afternoon and simple enough to explain to someone who doesn't do this for a living. The output is a centroid per cluster, which is a genuine summary you can use: the average customer in each segment, the prototype of each group. On data that actually falls into round, roughly equal blobs, it's hard to beat and pointless to complicate.

The bad side is the shape of its assumptions. k-means draws straight boundaries halfway between centroids and calls every cluster a round ball of similar size, because squared Euclidean distance is all it knows. Hand it two long stripes, or a small dense group next to a big sparse one, or clusters shaped like crescents, and it carves them wrong with total confidence. It needs you to pick k in advance, and it has no opinion about whether your k was any good. And it's sensitive to where the centroids start — a careless seed can settle into a bad local minimum and hand you a clustering that's stable, plausible, and wrong. Most of those weaknesses have a fix, and the biggest one, the seeding, we build in directly.

The data

Four Gaussian blobs from scikit-learn's make_blobs, 150 points in 2-D, with a fixed random seed so the snapshot is stable. Two features so we can put the whole thing on a page and actually see it; four centers spaced far enough apart that the right answer is obvious to a human, which is exactly what you want when you're checking whether the algorithm agrees. The generator also hands back a true blob id for each point. We commit those, but we treat them as a sealed envelope — the coordinates go into k-means, the ids come out only at grading time to compute how well the unsupervised result matched the structure that was really there.

Build it, one function at a time

Seven short functions. Each one is a step of Lloyd's algorithm or a piece of its bookkeeping, and they stack in the order you'd actually write them: the distance kernel first, then the two steps that use it, then the objective, then the seeding, then the loop that ties them together.

Everything rests on one distance computation. Given the points and the current centroids, we want the squared distance from every point to every centroid at once — a full grid, no loops:

def distances(X, centroids):
    """Squared Euclidean distance from every point to every centroid.

    X is (N, D), centroids is (k, D); the result is (N, k), where entry
    (i, j) is ||x_i - mu_j||^2. Broadcasting does the whole thing at once —
    no loop over points, no loop over centroids.
    """
    diff = X[:, None, :] - centroids[None, :, :]   # (N, k, D)
    return (diff ** 2).sum(axis=2)                 # (N, k)

That (N, k) grid is the workhorse. The assignment step is one line on top of it: for each point, which centroid is closest.

def assign(X, centroids):
    """Assignment step: label each point with its nearest centroid.

    Returns an (N,) array of cluster indices in [0, k). argmin over the
    distance matrix picks the closest centroid for every point at once.
    """
    return np.argmin(distances(X, centroids), axis=1)

argmin down the centroid axis turns the distance grid into a cluster label per point. That's the first half of a Lloyd iteration. The second half moves each centroid to the mean of the points that just chose it:

def update(X, labels, centroids):
    """Update step: move each centroid to the mean of its assigned points.

    k is read off the current centroids. A cluster that ended up empty keeps
    its old position — there is no mean to move to — so the loop never divides
    by zero and never loses a centroid.
    """
    new = centroids.copy()
    for j in range(len(centroids)):
        pts = X[labels == j]
        if len(pts):
            new[j] = pts.mean(axis=0)
    return new

The only wrinkle is an empty cluster — a centroid that no point picked has no mean to move to. We leave it where it is rather than divide by zero, which keeps k centroids alive for the whole run. Now the objective the two steps are quietly minimizing, so we can watch it fall:

def inertia(X, labels, centroids):
    """Within-cluster sum of squares — the objective k-means minimizes.

    Sum over points of the squared distance to the point's OWN centroid.
    Every Lloyd iteration drives this down; when it can't fall further, the
    algorithm has converged.
    """
    d2 = distances(X, centroids)
    return float(d2[np.arange(len(X)), labels].sum())

That's inertia straight from the math: pull each point's distance to its own centroid out of the grid and sum. Every iteration should make this number smaller or leave it flat, never larger — if it ever climbs, there's a bug.

Now seeding, and this is where k-means is won or lost. The naive way is to pick k random points and call them the starting centroids:

def init_random(X, k, rng):
    """The naive seed: k distinct data points picked uniformly at random.

    Simple, and sometimes fine — but nothing stops it from dropping two seeds
    inside the same blob and none in another, which is exactly the bad start
    that makes k-means take longer or converge to a worse answer.
    """
    idx = rng.choice(len(X), size=k, replace=False)
    return X[idx].copy()

It works often enough to be tempting, and it's what we'll use for the animation precisely because it's a little clumsy — it gives the centroids somewhere to walk from. But it can drop two seeds inside one blob and none in another, and from there k-means happily converges to a wrong-but-stable answer. k-means++ fixes it by spreading the seeds out on purpose: first center random, then each next center chosen with probability proportional to its squared distance from the centers already picked, so far-away regions are far likelier to get one:

def init_centroids(X, k, rng):
    """k-means++ seeding: spread the first centroids out on purpose.

    Pick the first centroid uniformly at random, then pick each next one with
    probability proportional to its squared distance from the centroids
    already chosen. Far-apart seeds are what save k-means from the bad local
    minima a naive random start can fall into.
    """
    first = int(rng.integers(len(X)))
    centroids = [X[first]]
    for _ in range(1, k):
        d2 = distances(X, np.array(centroids)).min(axis=1)
        probs = d2 / d2.sum()
        nxt = int(rng.choice(len(X), p=probs))
        centroids.append(X[nxt])
    return np.array(centroids)

That one change is the difference between an algorithm you have to babysit with restarts and one you can trust on the first try. Last, the loop that puts it all together — assign, update, repeat until the labels stop changing:

def kmeans(X, k, rng, max_iter=100, init=init_centroids):
    """Lloyd's algorithm: assign, update, repeat until labels stop changing.

    `init` chooses the seeding — k-means++ by default, or pass init_random for
    the naive version. Returns the final labels, centroids, and inertia, plus
    the full per-iteration history: a list of {centroids, labels, inertia}
    snapshots the chapter replays frame by frame so you watch the centroids move.
    """
    centroids = init(X, k, rng)
    labels = assign(X, centroids)
    history = [_snapshot(X, labels, centroids)]
    for _ in range(max_iter):
        centroids = update(X, labels, centroids)   # move to the means
        new_labels = assign(X, centroids)          # re-assign to new centroids
        history.append(_snapshot(X, new_labels, centroids))
        if np.array_equal(new_labels, labels):     # nothing moved — done
            labels = new_labels
            break
        labels = new_labels
    return labels, centroids, inertia(X, labels, centroids), history

It records a snapshot every iteration — centroids, labels, inertia — which is what the animation below replays. Convergence is the plainest possible test: if an assignment pass returns the exact labels we already had, no point can move again, and we're done.

Watch it work

This is the whole reason 2-D data is worth using. Below is a real run of the kmeans function above, one frame per iteration, seeded from random points so you can see it work rather than snap straight to the answer. Points are colored by their current cluster. The four diamonds are the centroids — watch them move. The dashed ring around each cluster reaches out to its farthest assigned point, so you can see each group's territory swell and shrink as the boundaries shift. The lower panel tracks inertia, the number the whole thing is driving down.

Press play. The centroids start scattered on arbitrary points and the coloring is a mess. Then each frame does one assign-and-average, and you watch the diamonds march into the middle of the blobs while the rings snap tight around them.

Watch the first two frames especially — that's where most of the work happens. The inertia drops from about 7,081 at the random start to roughly 2,456 after the first move and then keeps falling, a steep cliff that flattens fast. By the last frame no point changes cluster, the diamonds have stopped, and the run has converged in six iterations to an inertia of 342. That flattening is the whole story of Lloyd's algorithm: enormous progress in the first step or two, then fine-tuning, then a full stop. Reset and run it again — because the seeding here is random, but the blobs are clean, it walks to the same four groups every time.

That last part is luck of the data, not a guarantee. On messier data a random seed can strand a centroid and freeze in a worse local minimum, which is the argument for k-means++ made visible: spread the seeds and the good ending stops being luck.

The full implementation

The whole file, no library, top to bottom. This is exactly what the animation ran:

"""K-means clustering (Lloyd's algorithm), built from scratch.

No labels this time. Group unlabeled 2-D points into k clusters by alternating
two steps: assign every point to its nearest centroid, then move every centroid
to the mean of the points now assigned to it. Repeat until the assignments stop
changing. Pure NumPy — no ML library anywhere in this file.

Every function below appears in the chapter one step at a time (the
`# region:` markers are what the book's include directives pull in).
"""

import numpy as np
import pandas as pd


# region: distances
def distances(X, centroids):
    """Squared Euclidean distance from every point to every centroid.

    X is (N, D), centroids is (k, D); the result is (N, k), where entry
    (i, j) is ||x_i - mu_j||^2. Broadcasting does the whole thing at once —
    no loop over points, no loop over centroids.
    """
    diff = X[:, None, :] - centroids[None, :, :]   # (N, k, D)
    return (diff ** 2).sum(axis=2)                 # (N, k)
# endregion


# region: assign
def assign(X, centroids):
    """Assignment step: label each point with its nearest centroid.

    Returns an (N,) array of cluster indices in [0, k). argmin over the
    distance matrix picks the closest centroid for every point at once.
    """
    return np.argmin(distances(X, centroids), axis=1)
# endregion


# region: update
def update(X, labels, centroids):
    """Update step: move each centroid to the mean of its assigned points.

    k is read off the current centroids. A cluster that ended up empty keeps
    its old position — there is no mean to move to — so the loop never divides
    by zero and never loses a centroid.
    """
    new = centroids.copy()
    for j in range(len(centroids)):
        pts = X[labels == j]
        if len(pts):
            new[j] = pts.mean(axis=0)
    return new
# endregion


# region: inertia
def inertia(X, labels, centroids):
    """Within-cluster sum of squares — the objective k-means minimizes.

    Sum over points of the squared distance to the point's OWN centroid.
    Every Lloyd iteration drives this down; when it can't fall further, the
    algorithm has converged.
    """
    d2 = distances(X, centroids)
    return float(d2[np.arange(len(X)), labels].sum())
# endregion


# region: init_random
def init_random(X, k, rng):
    """The naive seed: k distinct data points picked uniformly at random.

    Simple, and sometimes fine — but nothing stops it from dropping two seeds
    inside the same blob and none in another, which is exactly the bad start
    that makes k-means take longer or converge to a worse answer.
    """
    idx = rng.choice(len(X), size=k, replace=False)
    return X[idx].copy()
# endregion


# region: init_centroids
def init_centroids(X, k, rng):
    """k-means++ seeding: spread the first centroids out on purpose.

    Pick the first centroid uniformly at random, then pick each next one with
    probability proportional to its squared distance from the centroids
    already chosen. Far-apart seeds are what save k-means from the bad local
    minima a naive random start can fall into.
    """
    first = int(rng.integers(len(X)))
    centroids = [X[first]]
    for _ in range(1, k):
        d2 = distances(X, np.array(centroids)).min(axis=1)
        probs = d2 / d2.sum()
        nxt = int(rng.choice(len(X), p=probs))
        centroids.append(X[nxt])
    return np.array(centroids)
# endregion


# region: kmeans
def kmeans(X, k, rng, max_iter=100, init=init_centroids):
    """Lloyd's algorithm: assign, update, repeat until labels stop changing.

    `init` chooses the seeding — k-means++ by default, or pass init_random for
    the naive version. Returns the final labels, centroids, and inertia, plus
    the full per-iteration history: a list of {centroids, labels, inertia}
    snapshots the chapter replays frame by frame so you watch the centroids move.
    """
    centroids = init(X, k, rng)
    labels = assign(X, centroids)
    history = [_snapshot(X, labels, centroids)]
    for _ in range(max_iter):
        centroids = update(X, labels, centroids)   # move to the means
        new_labels = assign(X, centroids)          # re-assign to new centroids
        history.append(_snapshot(X, new_labels, centroids))
        if np.array_equal(new_labels, labels):     # nothing moved — done
            labels = new_labels
            break
        labels = new_labels
    return labels, centroids, inertia(X, labels, centroids), history
# endregion


def _snapshot(X, labels, centroids):
    """One frame of the run: the centroids, the assignment to them, its cost."""
    return {
        "centroids": centroids.copy(),
        "labels": labels.copy(),
        "inertia": inertia(X, labels, centroids),
    }


def load_data(path="../data/blobs.csv"):
    """make_blobs snapshot: 150 2-D points (x, y) plus the true blob id.

    The blob column is ground truth we NEVER train on — k-means sees only the
    coordinates. It exists so we can score the clustering afterwards.
    """
    return pd.read_csv(path)

The library version

Nobody ships a hand-rolled k-means, and once you've built it you don't need to. scikit-learn's KMeans is the same Lloyd loop with the sharp edges filed off — k-means++ seeding by default, several random restarts so one unlucky seed can't decide the answer, and a C inner loop. It minimizes the identical inertia, so on clean data it should land where we did:

def sklearn_kmeans(X, k, seed=0):
    """Cluster X into k groups with scikit-learn. Returns labels, the cluster
    centers, and the final inertia — the same three things our kmeans returns."""
    km = KMeans(n_clusters=k, init="k-means++", n_init=10, random_state=seed)
    labels = km.fit_predict(X)
    return labels, km.cluster_centers_, float(km.inertia_)

The library also makes the k question cheap to answer. Inertia always falls as you add clusters — more centroids, shorter distances, all the way down to one point per cluster and zero inertia — so you can't just minimize it. What you look for instead is the elbow: the k where the curve stops dropping like a rock and goes nearly flat, the point of diminishing returns. Fit KMeans across a range of k and plot each converged inertia:

def elbow(X, ks):
    """Fit KMeans for a range of k and record each converged inertia.

    Inertia only ever falls as k grows, so you don't pick the minimum — you
    look for the 'elbow', the k where the curve stops dropping steeply and
    adding clusters buys you almost nothing.
    """
    out = []
    for k in ks:
        km = KMeans(n_clusters=k, n_init=10, random_state=0)
        km.fit(X)
        out.append((int(k), float(km.inertia_)))
    return out

The bend is about as clean as this ever gets. Inertia falls off a cliff from 7,881 at k=1 to 921.8 at k=3, drops once more to 342.2 at k=4, and then the curve goes nearly flat — k=5 only buys you down to 301.7, and each cluster after that shaves off less. The orange dot at k=4 is the elbow, and it's the right answer because we built the data with four blobs. In the wild the elbow is often mushier than this and you lean on it as a hint, not a verdict.

Scratch versus library

Same data, same k, both run to convergence — our from-scratch k-means++ against sklearn's KMeans. The bars are the final inertia:

They're identical — 342.2 and 342.2. That's not a coincidence and it's not a tie that needed rounding to hide; on four clean, well-separated blobs there's one obvious clustering, its inertia is 342.2, and any correct implementation that reaches the global optimum lands there. Our hand-written loop found it in two iterations from a k-means++ seed; sklearn found it too. When the answer is this unambiguous, the value of the library isn't a better number, it's the n_init restarts and the C loop that keep you at that number on data where the right clustering isn't sitting in plain sight.

Now open the sealed envelope. We never trained on the true blob ids, but we can grade against them with the adjusted Rand index, which scores how well two labelings agree after correcting for chance — 1.0 is a perfect match, 0 is random. Both our clustering and sklearn's score 1.0: every point landed in the group it was really drawn from. The unsupervised algorithm recovered the exact structure the data was built with, without ever being shown a single label.

Takeaways

This is the chapter where the course crosses a line. Everything before it was supervised — a label column, a target to match, an accuracy to report. k-means has none of that. It learns the shape of the data from the data, and the only reason we could put a number on it here is that we happened to know the truth and kept it hidden until the end. On real unlabeled data you don't get that envelope. You get inertia, an elbow, and your own judgment about whether the groups it found mean anything. Clustering moves the hard question from "is the model right" to "is the model useful", and no metric answers that for you.

Reach for k-means first when you need to segment something and you believe the groups are roughly round and roughly balanced — it's fast, it's simple, it gives you a usable centroid per cluster, and with k-means++ seeding it's reliable enough to trust on the first run. Reach past it when those assumptions break: when you don't know k and don't want to commit to one, when clusters nest inside each other, or when a cluster is a long streak or a crescent instead of a ball. The next two chapters are exactly those escapes. Hierarchical clustering refuses to fix k up front and hands you the whole tree of groupings to cut where you like. Gaussian mixtures keep the centroids but let clusters be stretched, tilted ellipses with soft edges, so a point can belong partly to two groups at once. Both are k-means with an assumption relaxed — which is the best possible reason to have built k-means by hand first.