k-means Clustering, in depth: the Lloyd algorithm, a hand trace, k-means++ seeding, choosing k and scaling up

The objective

Given n points x₁..x_n in d dimensions and a number k, k-means looks for k centres c₁..c_k and an assignment of each point to one centre that minimises the within-cluster sum of squared distances, usually called SSE or inertia:

SSE = sum over points i of ||x_i - c_a(i)||², where a(i) is the cluster of point i.

Two facts make the problem tractable to attack even though it is hard to solve exactly. First, if the centres are fixed, the best assignment is obvious: send every point to its nearest centre. Second, if the assignment is fixed, the best centre for each cluster is the mean of its points, because the mean is the unique minimiser of a sum of squared Euclidean distances (set the gradient 2 times the sum of (c - x_i) to zero). Each half of the problem is easy; the joint problem is NP-hard in general, even for k = 2 when the dimension is part of the input. That is why every practical algorithm is a local search, and why the starting point matters.

Squared distance is not a detail. It is what makes the mean optimal and it is why a single far outlier can drag a centre: its contribution grows with the square of its distance. If you need a centre that must be an actual data point, or a cost that is robust to outliers, you want k-medoids or k-medians instead.

Lloyd iterations and why they stop

Lloyd's algorithm alternates the two easy halves until nothing changes:

initialise centres c_1..c_k
repeat:
    # assignment step: nearest centre for every point
    for each point i: a(i) = argmin_j ||x_i - c_j||^2
    # update step: each centre moves to the mean of its points
    for each cluster j with at least one point: c_j = mean of {x_i : a(i) = j}
until no assignment changes (or SSE improves by less than tol)

Why it terminates. The assignment step can only lower SSE or leave it unchanged, since each point moves to a centre at least as close. The update step can only lower SSE or leave it unchanged, since the mean is the optimal centre for a fixed cluster. So SSE never increases. There are finitely many ways to partition n points into k groups, and with a consistent tie-breaking rule the algorithm cannot revisit a partition without having stopped, so it reaches a fixed point in finitely many steps. That fixed point is a local minimum with respect to these moves, not necessarily the global one. In the worst case the number of iterations can be exponential in n, but on real data a few dozen iterations are typical.

Cost per iteration. The assignment step computes n times k distances in d dimensions, O(nkd), and dominates. The update step is O(nd). Memory is O(nd + kd) plus the assignments.

A hand trace on seven points

Take seven points in the plane: A(1, 1), B(1.5, 2), C(3, 4), D(5, 7), E(3.5, 5), F(4.5, 5) and G(3.5, 4.5), with k = 2 and initial centres placed on A and D.

1. Assign. Squared distances to (1, 1) and (5, 7): A 0 vs 52, B 1.25 vs 37.25, C 13 vs 13, D 52 vs 0, E 22.25 vs 6.25, F 28.25 vs 4.25, G 18.5 vs 8.5. C is an exact tie; breaking ties toward the lower index puts it in cluster 1. Clusters {A, B, C} and {D, E, F, G}; SSE = 33.25.
2. Update. Means: (1.833, 2.333) and (4.125, 5.375). With the same clusters SSE drops to 12.21.
3. Assign. C is now at squared distance 4.14 from the first centre and 3.16 from the second, so it switches. SSE with the new clusters, before moving centres: 11.23.
4. Update. {A, B} has mean (1.25, 1.5) and {C, D, E, F, G} has mean (3.9, 5.1). SSE = 8.525.
5. Assign. No point changes cluster, so the algorithm stops.

The sequence 33.25, 12.21, 11.23, 8.525 is monotone, as the proof says it must be. Brute force over all partitions confirms 8.525 is the global optimum for k = 2 here. Note how much the tie mattered early on: a library that broke the tie the other way would take a different path, which is one reason two implementations can disagree on the same data and seed.

Initialisation and k-means++

Lloyd's algorithm only finds a local minimum, so initialisation decides quality. Run the same seven points with k = 3 and the bad start A, B, C: the result is clusters {A}, {B} and the other five, with SSE 7.90. Starting from A, D, E instead gives {A, B}, {D} and {C, E, F, G} with SSE 2.50, the global optimum. Same data, same algorithm, three times the cost.

k-means++ (Arthur and Vassilvitskii, 2007) fixes the starting point cheaply. Pick the first centre uniformly at random. Then pick each next centre from the data with probability proportional to D(x)², the squared distance from x to the nearest centre chosen so far. From the first centre A in our example, the D(x)² values are B 1.25, C 13, D 52, E 22.25, F 28.25 and G 18.5, totalling 135.25, so D is picked with probability 38% and B with under 1%. Far, uncovered regions get centres. The expected SSE of the seeding alone is within O(log k) of optimal, and Lloyd iterations only improve it. Seeding costs k passes over the data, so for very large n use a parallel variant that samples several centres per pass.

Even with k-means++, run several independent initialisations and keep the lowest final SSE. Five to ten restarts are common; the cost is linear in the number of restarts and they parallelise perfectly.

A NumPy implementation

A compact, readable implementation with k-means++ seeding, vectorised distances, empty-cluster repair and restarts:

import numpy as np

def kmeans_pp_init(X, k, rng):
    n = X.shape[0]
    centers = [X[rng.integers(n)]]
    d2 = ((X - centers[0]) ** 2).sum(axis=1)
    for _ in range(1, k):
        idx = rng.choice(n, p=d2 / d2.sum())
        centers.append(X[idx])
        d2 = np.minimum(d2, ((X - X[idx]) ** 2).sum(axis=1))
    return np.array(centers, dtype=float)

def lloyd(X, C, max_iter=300, tol=1e-6):
    for _ in range(max_iter):
        # ||x - c||^2 = ||x||^2 - 2 x.c + ||c||^2, computed for all pairs at once
        d2 = (X**2).sum(1)[:, None] - 2 * X @ C.T + (C**2).sum(1)[None, :]
        labels = d2.argmin(axis=1)
        newC = C.copy()
        for j in range(len(C)):
            members = X[labels == j]
            if len(members):
                newC[j] = members.mean(axis=0)
            else:  # empty cluster: re-seed at the point worst served
                newC[j] = X[d2[np.arange(len(X)), labels].argmax()]
        shift = ((newC - C) ** 2).sum()
        C = newC
        if shift < tol:
            break
    d2 = ((X[:, None, :] - C[None, :, :]) ** 2).sum(-1)
    labels = d2.argmin(1)
    return C, labels, d2[np.arange(len(X)), labels].sum()

def kmeans(X, k, n_init=8, seed=0):
    rng = np.random.default_rng(seed)
    runs = [lloyd(X, kmeans_pp_init(X, k, rng)) for _ in range(n_init)]
    return min(runs, key=lambda r: r[2])   # (centers, labels, sse)

Two production notes. The expanded-square distance trick turns the assignment step into one matrix multiply, which is why GPU implementations are fast, but it can go slightly negative from rounding, so clamp at zero if you take square roots. The final pass recomputes distances directly to report an exact SSE.

Choosing k

k-means never tells you k; SSE always falls as k grows and reaches zero at k = n. For our points the optimal SSE is 37.07 at k = 1, 8.53 at k = 2 and 2.50 at k = 3. Common ways to choose:

• Elbow. Plot SSE against k and look for the bend where extra clusters stop paying. Here the drop from 1 to 2 removes 77% of SSE and the drop from 2 to 3 removes another 71% of what is left, so the elbow is ambiguous, which is common on real data too.
• Silhouette. For each point, compare its mean distance to its own cluster (a) with its mean distance to the nearest other cluster (b): s = (b - a) / max(a, b). Average s near 1 means tight, separated clusters; pick the k that maximises it.
• Downstream metric. Usually the honest answer. For an IVF index k trades recall against probe cost, for a palette k is set by the format, and for data curation k is chosen by how duplicates are reviewed. Choose k with the metric the clusters serve.

Scaling to millions of points

At millions of points and thousands of centres, the O(nkd) assignment step dominates. The main tools:

• Mini-batch k-means. Update centres from small random batches with a per-centre learning rate of 1/count. It is much faster per epoch and slightly worse in final SSE; good for streaming data and huge n.
• Triangle-inequality pruning. Elkan and Hamerly variants keep bounds on each point's distance to its centre and to the others, and skip distance computations that cannot change the assignment. Results are identical to Lloyd, often with far fewer distance evaluations at large k.
• Train on a sample, assign everything. IVF indexes train the coarse quantizer on a sample of a few hundred points per centre, then assign the full collection once. See IVF-PQ vector index architecture and the codebook maths in PQ Math.
• GPU batching. Assignment is a matrix multiply plus argmin; chunk X so the n by k distance block fits in memory.

Failure modes

k-means encodes assumptions, and data that violates them produces confident nonsense:

• Unscaled features. A feature measured in thousands dominates one measured in units. Standardise or choose a meaningful scale first.
• Non-convex shapes. Concentric rings or interleaved moons get cut by straight boundaries, because every k-means boundary is a perpendicular bisector between centres. Use density-based or graph-based methods, such as Label Propagation on a neighbour graph.
• Unequal sizes and spreads. SSE favours splitting a big diffuse cluster over keeping a small tight one separate. A Gaussian mixture with per-cluster covariances handles this.
• Outliers. Squared distance lets one far point pull a centre or claim a cluster of its own. Trim outliers or use k-medoids.
• High-dimensional embeddings. Euclidean distances concentrate; normalise vectors and use spherical k-means (cosine) when direction is what matters.
• Instability. Different seeds give different labels. Fix seeds for reproducibility and never compare cluster IDs across runs without matching them first.

What to do next

1. Implement Lloyd by hand on the seven-point example and reproduce SSE 33.25, 12.21, 11.23, 8.525.
2. Add k-means++ seeding and restarts; compare final SSE against random seeding on your own data.
3. Standardise features, then choose k with silhouette and with the metric your downstream task uses.
4. Check cluster shapes visually (PCA or UMAP); switch to a mixture or density method if they are not blobs.
5. For large n, train on a sample with mini-batch or Elkan k-means and assign the full set once.
6. Record seed, k, init and tolerance with every clustering you ship, so results can be reproduced.