What it actually does
Summarize the data with k representative points. Start with k guesses, assign every observation to its nearest guess, then move each guess to the centre of what it just captured, and repeat until nothing moves. The objective it descends is the total squared distance from points to their own centre, and the two steps alternate between the two things that objective depends on — which assignment, and which centres. What makes it so widely used is that both steps are trivial and it converges quickly. What makes it so widely misused is that it always returns exactly k clusters, whether or not the data has any.
- Clusters are roughly spherical and comparable in size, which is what squared Euclidean distance encodes — elongated or unequal clusters are split and merged incorrectly
- Euclidean distance is meaningful for the features, so they must be scaled comparably and not too high-dimensional
- The number of clusters is known or can be chosen externally, since the objective always prefers more
- Clusters actually exist — the algorithm returns K partitions of uniform noise just as confidently as of genuinely clustered data
- The mean is a sensible representative, which fails when outliers are present and is what k-medoids exists to fix
Loss function
Add up how far every point is from the centre of its own cluster, squared, and make that total as small as possible. Two consequences hide in the notation. Squared Euclidean distance is what makes the mean the right centre — under any other distance the optimal representative is not a mean — so the algorithm and the metric are not separable choices. And J decreases monotonically as K grows, reaching zero when every point is its own cluster, which means the objective offers no opinion whatsoever on how many clusters there are. Every method for picking K is external to the model.
- the set of points assigned to cluster k — a hard partition, with no point split between clusters
- the centroid: the mean of the points assigned to k, which is what squared error makes optimal
- the number of clusters, which the objective cannot choose — J falls monotonically as K rises
- squared Euclidean distance, and therefore an assumption of spherical, comparably-sized clusters
How it is minimized
Lloyd’s algorithm — alternating minimization, with k-means++ seeding
- the two alternating steps: fix the centres and reassign, then fix the assignment and recentre
- nearest-centroid assignment — the step that costs O(n·K·d) and dominates the runtime
- cluster size; an empty cluster makes this zero, which every implementation must handle explicitly
- the mean, which minimizes squared distance exactly — this is why the metric and the update are locked together
Minimizing J over both the assignment and the centres jointly is NP-hard, so Lloyd’s algorithm minimizes over one while holding the other fixed and alternates. Each step is exactly optimal for its own subproblem, so J cannot increase and the algorithm terminates in finitely many steps — there are only finitely many partitions, and none is visited twice. That is a convergence guarantee about the algorithm, not about the answer: the fixed point it reaches depends entirely on where it started. This is the same alternating-minimization pattern as EM in a Gaussian mixture, with hard assignments in place of responsibilities, and it inherits the same non-convexity. The one intervention that reliably matters is initialization: k-means++ seeds centres far apart with probability proportional to squared distance, and it turns a method that regularly converges somewhere embarrassing into one that usually does not.
K
The number of clusters, and the only decision that really matters. Not learnable from J, so it comes from the elbow, silhouette, gap statistic, or the application (chosen externally; every internal criterion is a heuristic and they routinely disagree)
init
k-means++ or random. This is the difference between a reliable method and a lottery, and there is no reason to use random (k-means++, always)
n_init
Restarts, keeping the run with the lowest J. Still worth several even with k-means++, since the surface is non-convex (10 with random init, 1 to 3 with k-means++)
max_iter / tol
Stop on iteration count or on centroid movement. Rarely binding — convergence is usually tens of iterations (300 iterations, tol 1e-4 on centroid shift)
algorithm
Lloyd, Elkan (triangle-inequality bounds that skip distance computations), or mini-batch. A cost decision; mini-batch trades a slightly worse optimum for a large speedup
Guaranteed to terminate at a local optimum in finitely many iterations, and that guarantee is weaker than it sounds — different seeds give materially different partitions, which is why restarts are part of the method. Three failure modes deserve names. Empty clusters: a centroid can capture nothing, its update divides by zero, and the implementation must reseed it, usually to the point furthest from any centre. Shape mismatch: squared Euclidean distance assumes spherical, comparably-sized clusters, so an elongated cluster gets split down the middle and a small cluster next to a large one gets absorbed. And the quietest one — the algorithm always returns K clusters, so it will happily partition uniform noise and report a tidy result, which no internal metric will flag as meaningless.
O(n·K·d) per iteration, dominated entirely by the assignment step, times typically tens of iterations. k-means++ seeding costs an extra O(n·K·d) up front and is worth it. Memory is O(n·d + K·d), with no n-by-n or n-by-K matrix needed if the assignment is computed in chunks. Mini-batch reduces the per-iteration cost to O(b·K·d) for batch size b, which is what makes it viable on data that does not fit in memory.
How, where, and why it is used
Time-Series Forecasting
AdaptedCluster series by shape rather than forecasting with it. Normalize each series, cluster the resulting curves, and either fit one model per cluster — pooling data across similar series — or use the cluster label as a feature. On thousands of sparse series this is often what makes modelling possible at all, since no individual series has enough history.
- Grouping SKUs or meters by demand profile so that a model can be fitted per group rather than per series
- Load-profile segmentation in energy, where customer classes are discovered rather than assumed
- Clustering learned or engineered series embeddings, where the features already encode shape
It is the cheapest way to exploit the fact that many series behave alike, and pooling across a cluster is frequently worth more than a better model on each series alone. The caveat is specific and important: Euclidean distance on raw series compares values timestep by timestep, so two identical shapes offset by one period look completely different. Dynamic time warping fixes that and is not compatible with a mean, which is why DTW clustering needs a different centroid definition entirely — it is a different algorithm, not a distance swap.
- Z-normalize each series before clustering, or the clusters simply recover magnitude rather than shape
- Cluster on engineered features or a learned embedding rather than raw points when series differ in length or phase
- Fit the clustering on the training window only; refitting it across the full history leaks future shape into the labels
- Pin the cluster ordering after fitting — label switching makes a cluster id meaningless across refits
Rolling-origin backtest of the pooled forecasts, not a clustering metric. A partition that improves silhouette and does not improve forecast error has described the data rather than helped it.
- Euclidean distance on unaligned series, where a one-period phase shift dominates any difference in shape
- Clusters that reorder between refits, silently changing which pooled model a series is scored by
- Choosing K by an internal criterion when the downstream forecast error is the criterion that matters
Anomaly Detection
AdaptedCluster the normal data and score each new point by its distance to the nearest centroid. Points far from every centre are candidates. A second, sharper signal is cluster size — a tiny cluster is often a group of anomalies that found each other rather than a genuine mode.
- Fast novelty screening at volumes where a density model would be too slow to refit
- Network intrusion baselines built on clustered normal traffic profiles
- A coarse first pass ahead of a more expensive detector, discarding the obviously typical
It is the crude version of what a Gaussian mixture does properly, and its virtue is entirely operational: assignment is a single nearest-centroid lookup, the model is K points, and both fit and score are far cheaper than a density model. Its weakness follows directly from the spherical equal-variance assumption — the score is a raw distance, so the same distance means something different next to a tight cluster than next to a diffuse one, and there is no per-cluster scale to correct with. When calibration matters, the mixture is the right model; when throughput matters, this is defensible.
- Standardize before clustering, since the distance is Euclidean and unscaled features dominate it
- Fit on a confirmed-clean window, or anomalies form their own cluster and become perfectly normal
- Normalize the distance by each cluster’s own spread before thresholding, which recovers part of what the spherical assumption discards
- Reduce dimension first, since Euclidean distance degrades well before the feature count looks alarming
Precision@k against confirmed incidents with thresholds set per cluster rather than globally, and PR-AUC rather than ROC-AUC. Compare against a mixture on the same data — if the mixture is much better, the spherical assumption is the thing costing you.
- A global distance threshold across clusters of different spread, which over-alarms on the tight ones and under-alarms on the diffuse ones
- Anomalies forming their own cluster during fitting and being scored as central members of it
- Contamination pulling a centroid toward the anomalies, which moves normal points away from their own centre
Optimization
AdaptedThe objective is NP-hard to minimize jointly, and Lloyd’s algorithm is the textbook alternating-minimization heuristic for it: optimize the assignment with the centres fixed, then the centres with the assignment fixed, and repeat. That is the same structure as EM, as coordinate descent in the Lasso, and as block-coordinate methods generally — this is the smallest example where the pattern is fully visible in a dozen lines.
- The canonical worked example of alternating minimization and why monotone descent does not imply a global optimum
- Vector quantization for compression, where the objective genuinely is the deliverable rather than a proxy
- Facility-location and warehouse-siting problems, whose continuous relaxation is exactly this objective
It is worth studying here because the gap between the problem and the algorithm is unusually explicit: the objective is NP-hard, the algorithm is a heuristic with a termination guarantee and no approximation guarantee, and k-means++ closes part of that gap with an actual bound — expected cost within O(log K) of optimal, purely from how the seeds are drawn. That is a rare, concrete demonstration that initialization can carry a theoretical guarantee rather than being folklore.
- Standardize so that the squared-distance objective weights every dimension comparably
- Where the application is genuinely facility location, encode real costs rather than assuming Euclidean distance is the objective
Compare the final inertia across many restarts — the spread between the best and worst run is a direct measurement of how non-convex the surface is on this data, and it is more informative than any single run’s value.
- Reading monotone descent as convergence to the optimum, when the guarantee is only that it stops
- Comparing inertia across different K, which always favours larger K and says nothing
- Using random initialization and attributing the resulting variance to the data rather than to the seeding
Computer Vision
AdaptedVector quantization. Colour quantization treats each pixel as a point in colour space and replaces it with its nearest of K centroids, which is image compression by clustering. Bag-of-visual-words does the same to local descriptors: cluster them to build a visual vocabulary, then represent an image by the histogram of which words it contains.
- Palette reduction and image compression, where the centroids are literally the palette
- Bag-of-visual-words pipelines, the standard image representation before learned features
- Superpixel segmentation, where a spatially-constrained variant groups adjacent similar pixels
- Building a codebook for a downstream retrieval index
It works here because the problem genuinely is quantization — find K representatives that reconstruct the data well — which is exactly the objective rather than a proxy for it. Colour space is also low-dimensional and roughly isotropic, so the spherical assumption is nearly true. For recognition it has been superseded: a visual vocabulary discards spatial arrangement entirely, which is precisely the information a convolution keeps.
- Work in a perceptually uniform colour space, or Euclidean distance does not correspond to visible difference
- Subsample pixels for fitting; a full-resolution image has far more points than the centroids need
- For visual words, L2-normalize descriptors first so the vocabulary is not dominated by high-contrast patches
Reconstruction error or a perceptual metric for quantization; downstream classification or retrieval accuracy for a vocabulary. Cluster quality metrics are not the deliverable in either case.
- Clustering in RGB and being surprised that the palette looks wrong, since RGB distance is not perceptual distance
- Vocabulary size chosen by a clustering criterion rather than by downstream accuracy
- Refitting the codebook without retraining what consumes it, which reindexes every visual word
Recommendation & Ranking
ViableTwo roles. Segmenting users or items by behaviour, so that a cold-start user can be served their segment’s popular items rather than a global default. And quantizing an embedding space to build the coarse index that approximate nearest-neighbour search uses — an inverted file index is literally k-means over the vectors, with the search restricted to the nearest few cells.
- User and item segmentation for cold start and for coarse targeting
- The coarse quantizer inside an IVF vector index, which is where most large-scale retrieval systems use it
- Reducing a catalogue to representative items for exploration or merchandising
As an index quantizer it is genuinely the standard, because the job is exactly vector quantization and the cost profile fits — fit once offline, then assignment is a single nearest-centroid lookup at query time. As a segmentation tool it is a reasonable coarse instrument and a poor personalization model: a hard assignment to one of twenty segments discards nearly everything a matrix factorization or two-tower model would use.
- L2-normalize embeddings before clustering when the retrieval metric is cosine, so the quantizer and the search agree
- Down-weight popularity before behavioural segmentation, or the clusters recover activity level rather than taste
- Choose the cell count from the index recall-latency curve, not from a clustering criterion
For an index, recall@k against exact search at a fixed latency budget — that curve is the only thing that matters. For segmentation, downstream engagement lift against a global-popularity baseline.
- Unbalanced cells in an IVF index, where one enormous cell destroys the latency guarantee the index exists to provide
- Treating segment membership as a personalization signal when it is a twenty-way bucket
- Refitting the quantizer without rebuilding the index, so cell assignments no longer match the stored lists
What it costs to run
O(n·K·d) per iteration over tens of iterations, so seconds to minutes on millions of points — and mini-batch removes even that constraint by working on samples. Restarts multiply it directly, which is the main reason a k-means fit is ever slow.
Assignment is a nearest-centroid lookup: K distance computations, microseconds, with a model of just K·d floats. That combination is why it is the quantizer of choice inside systems that must decide something per request.
Cheap enough to refit often, and the reason not to is stability rather than cost. Every refit produces a different labelling of the same structure, so anything keyed to a cluster id must be remapped or re-approved, and downstream indexes built on the old centroids have to be rebuilt with them.
- Track inertia per cluster over time — a cluster whose average distance is growing is losing coherence, and it moves before any downstream metric does
- Watch cluster sizes; a cluster collapsing toward empty or absorbing everything is a structural change worth investigating rather than absorbing
- Monitor the share of points far from every centroid, which is the population arriving that the current partition does not describe
- Compare centroid positions across refits after matching them up, since raw cluster indices are not comparable
- Cluster ids are arbitrary and change on every refit — anything downstream must key on a canonicalized ordering or on the centroid itself, never on the raw index
- The scaler is part of the model: refitting it separately changes every distance and therefore every assignment
- Empty clusters must be handled explicitly, or the centroid update divides by zero and the fit produces NaN centres that silently swallow every point
- Determinism requires pinning the seed, since both k-means++ and random init are stochastic and two runs give different partitions
- A new point far from every centroid is still assigned to the nearest one with no indication that it does not belong — the distance must be checked separately if that matters
Pros and cons, in context
Fast, simple, and scales to very large datasets
Linear per iteration in every dimension of the problem, with a mini-batch variant for data that does not fit in memory. This is the reason it is the default clustering algorithm despite every limitation below.
The model is K points, and assignment is one lookup
Decisive when clustering has to happen inside a request path or an index — an IVF vector index is built on exactly this property. Irrelevant when the clustering is an offline analysis.
Guaranteed to terminate, with a monotone objective
Convergence is finite and checkable, and the inertia gives a directly comparable quality number across restarts. The guarantee is about stopping, not about the answer, which is the part most often overstated.
k-means++ gives initialization a real guarantee
An expected cost within a logarithmic factor of optimal, purely from how the seeds are drawn — a rare case where a practical heuristic carries an actual bound. It does not remove the need for restarts.
It always returns K clusters, whether or not any exist
The most under-acknowledged failure here: run it on uniform noise and it produces a clean, plausible partition that no internal metric flags. Any clustering result needs a separate argument that structure was there to find.
Assumes spherical, equally-sized clusters
An elongated cluster is split down the middle and a small cluster beside a large one is absorbed. This is precisely the restriction a Gaussian mixture lifts, which is why the two entries belong side by side.
K must be chosen, and the objective cannot help
Inertia falls monotonically with K, so every selection method — elbow, silhouette, gap — is external and they routinely disagree. On real data the honest answer is often that no K is clearly right.
Sensitive to scaling, outliers, and initialization
A single extreme point drags a centroid, an unscaled feature decides the partition, and a bad seed reaches a different local optimum. All three are manageable and all three are routinely unmanaged.
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""k-means by Lloyd's algorithm - the two steps, transcribed.
Assign every point to its nearest centroid, then move every centroid to the
mean of what it captured. Each step is exactly optimal for its own subproblem
with the other held fixed, which is why the objective can only fall.
"""
import random
def squared_distance(a, b):
total = 0.0
for j in range(len(a)):
difference = a[j] - b[j]
total += difference * difference
return total
def fit(X, k, iterations=300, tol=1e-4, seed=0):
n = len(X)
d = len(X[0])
rng = random.Random(seed)
# Random initialization: this is the version k-means++ replaces, and the
# difference between them is the single largest practical improvement
# available to this algorithm.
centroids = [list(X[i]) for i in rng.sample(range(n), k)]
assignment = [0] * n
for _ in range(iterations):
# ---- assignment step: argmin over centroids ----------------------
for i in range(n):
best_cluster = 0
best_distance = squared_distance(X[i], centroids[0])
for cluster in range(1, k):
distance = squared_distance(X[i], centroids[cluster])
if distance < best_distance:
best_cluster = cluster
best_distance = distance
assignment[i] = best_cluster
# ---- update step: centroid becomes the mean of its members -------
totals = [[0.0] * d for _ in range(k)]
counts = [0] * k
for i in range(n):
cluster = assignment[i]
counts[cluster] += 1
for j in range(d):
totals[cluster][j] += X[i][j]
shift = 0.0
for cluster in range(k):
if counts[cluster] == 0:
# An empty cluster divides by zero. Reseeding to a random
# point is the crude fix; reseeding to the point furthest
# from any centroid is the better one.
centroids[cluster] = list(X[rng.randrange(n)])
continue
for j in range(d):
updated = totals[cluster][j] / counts[cluster]
shift += abs(updated - centroids[cluster][j])
centroids[cluster][j] = updated
if shift < tol:
break
return centroids, assignment
def inertia(X, centroids, assignment):
"""J = sum over points of the squared distance to their own centroid."""
total = 0.0
for i in range(len(X)):
total += squared_distance(X[i], centroids[assignment[i]])
return totalO(n*k*d) per iteration in interpreter loops, with random initialization that regularly converges somewhere poor.
How this language actually wants it written.
"""k-means - typed, k-means++ seeded, vectorized assignment."""
from dataclasses import dataclass
import numpy as np
from numpy.typing import NDArray
Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]
Labels = NDArray[np.int64]
@dataclass(frozen=True)
class KMeansModel:
"""A fitted model: K centroids, and the inertia that produced them."""
centroids: Matrix # (K, d)
inertia: float
n_iter: int
converged: bool
def predict(self, X: Matrix) -> Labels:
if X.ndim != 2 or X.shape[1] != self.centroids.shape[1]:
raise ValueError(
f"expected (n, {self.centroids.shape[1]}) input, got {X.shape}"
)
return _assign(X, self.centroids)[0]
def distance_to_centroid(self, X: Matrix) -> Vector:
"""The anomaly score. Raw distance, so it is NOT comparable across
clusters of different spread - normalize per cluster before
thresholding, which is what the spherical assumption costs you."""
return np.sqrt(_assign(X, self.centroids)[1])
def _assign(X: Matrix, centroids: Matrix) -> tuple[Labels, Vector]:
"""Nearest centroid and its squared distance, for every point at once."""
squared = ((X[:, None, :] - centroids[None, :, :]) ** 2).sum(axis=2)
labels = squared.argmin(axis=1)
return labels, squared[np.arange(X.shape[0]), labels]
def _kmeans_plus_plus(X: Matrix, k: int, rng: np.random.Generator) -> Matrix:
"""Seed centroids far apart, with probability proportional to D^2.
This is the single most consequential change available to k-means: random
seeding regularly lands in a poor local optimum, while this carries an
actual guarantee - expected cost within O(log k) of optimal, purely from
how the seeds are drawn.
"""
centroids = np.empty((k, X.shape[1]), dtype=np.float64)
centroids[0] = X[rng.integers(X.shape[0])]
closest = ((X - centroids[0]) ** 2).sum(axis=1)
for index in range(1, k):
total = closest.sum()
if total <= 0.0: # every point already a centroid
centroids[index] = X[rng.integers(X.shape[0])]
continue
chosen = int(rng.choice(X.shape[0], p=closest / total))
centroids[index] = X[chosen]
# Keep only the distance to the NEAREST centroid so far.
np.minimum(closest, ((X - centroids[index]) ** 2).sum(axis=1), out=closest)
return centroids
def fit(
X: Matrix,
n_clusters: int,
max_iter: int = 300,
tol: float = 1e-4,
n_init: int = 3,
seed: int = 0,
) -> KMeansModel:
"""Fit by Lloyd's algorithm. Raises ValueError on malformed input."""
if X.ndim != 2:
raise ValueError(f"X must be 2-D, got shape {X.shape}")
if not 1 <= n_clusters <= X.shape[0]:
raise ValueError(f"n_clusters must lie in [1, {X.shape[0]}], got {n_clusters}")
if n_init < 1:
raise ValueError(f"n_init must be at least 1, got {n_init}")
rng = np.random.default_rng(seed)
best: KMeansModel | None = None
# The surface is non-convex, so restarts are part of the method rather
# than a precaution. The spread across runs measures how non-convex it is.
for _ in range(n_init):
centroids = _kmeans_plus_plus(X, n_clusters, rng)
converged = False
iteration = 0
for iteration in range(1, max_iter + 1):
labels, _ = _assign(X, centroids)
updated = np.zeros_like(centroids)
counts = np.bincount(labels, minlength=n_clusters)
np.add.at(updated, labels, X)
empty = counts == 0
if empty.any():
# Reseed an empty cluster to the point furthest from any
# centroid, which is where a cluster is most likely missing.
_, distances = _assign(X, centroids)
furthest = np.argsort(distances)[-int(empty.sum()):]
updated[empty] = X[furthest]
counts[empty] = 1
updated /= counts[:, None]
shift = float(np.abs(updated - centroids).max())
centroids = updated
if shift < tol:
converged = True
break
labels, distances = _assign(X, centroids)
candidate = KMeansModel(
centroids=centroids,
inertia=float(distances.sum()),
n_iter=iteration,
converged=converged,
)
if best is None or candidate.inertia < best.inertia:
best = candidate
assert best is not None
return bestWhat changed: The change that matters most is not the vectorization: it is k-means++ seeding. Random initialization regularly converges to a poor local optimum, and seeding proportional to squared distance carries an actual guarantee — expected cost within a logarithmic factor of optimal — which turns a lottery into a reliable method. Restarts become explicit for the same reason. On top of that, the assignment step becomes a single broadcast reduction, empty clusters are reseeded to the furthest point rather than to a random one, and the fitted model is a frozen dataclass that reports its own inertia and whether it actually converged.
O(n*k*d) per iteration, but the broadcast materializes an (n, k, d) intermediate — which is the memory ceiling, not the arithmetic.
Faster, and the justification for each change.
"""k-means - squared-norm expansion, chunked assignment, in-place updates."""
import numpy as np
from numpy.typing import NDArray
Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]
Labels = NDArray[np.int64]
class ChunkedKMeans:
"""Assignment as a matrix product, over bounded chunks.
The broadcast form of the assignment step materializes an (n, k, d)
intermediate, which at a million points and a hundred clusters in fifty
dimensions is 40 GB. Expanding the square removes it entirely:
||x - c||^2 = ||x||^2 - 2 x . c + ||c||^2
Only the cross term touches both operands, and it is a GEMM. The ||x||^2
term is constant within a row and does not affect the argmin, so it is
never computed at all.
"""
def __init__(self, n_clusters: int, chunk: int = 8192) -> None:
if n_clusters < 1:
raise ValueError(f"n_clusters must be at least 1, got {n_clusters}")
self._k = n_clusters
self._chunk = chunk
def fit(self, X: Matrix, max_iter: int = 300, tol: float = 1e-4, seed: int = 0) -> Matrix:
design = np.ascontiguousarray(X, dtype=np.float64)
n, d = design.shape
rng = np.random.default_rng(seed)
centroids = np.ascontiguousarray(
design[rng.choice(n, size=self._k, replace=False)], dtype=np.float64
)
# Sufficient statistics, allocated once and reused every iteration.
totals = np.empty((self._k, d), dtype=np.float64)
counts = np.empty(self._k, dtype=np.int64)
labels = np.empty(n, dtype=np.int64)
for _ in range(max_iter):
centroid_norms = np.einsum("ij,ij->i", centroids, centroids)
totals.fill(0.0)
counts.fill(0)
# Chunking bounds the (chunk, k) score block rather than (n, k),
# which is what makes a large n possible in one pass.
for start in range(0, n, self._chunk):
stop = min(start + self._chunk, n)
block = design[start:stop]
scores = block @ centroids.T # the GEMM
scores *= -2.0 # in place
scores += centroid_norms # ||x||^2 omitted: constant per row
block_labels = scores.argmin(axis=1)
labels[start:stop] = block_labels
# Accumulate sufficient statistics per chunk; the full
# assignment matrix is never held.
np.add.at(totals, block_labels, block)
counts += np.bincount(block_labels, minlength=self._k)
empty = counts == 0
if empty.any():
totals[empty] = design[rng.choice(n, size=int(empty.sum()))]
counts[empty] = 1
updated = totals / counts[:, None]
shift = float(np.abs(updated - centroids).max())
np.copyto(centroids, updated)
if shift < tol:
break
return centroids
def predict(self, X: Matrix, centroids: Matrix) -> Labels:
centroid_norms = np.einsum("ij,ij->i", centroids, centroids)
out = np.empty(X.shape[0], dtype=np.int64)
for start in range(0, X.shape[0], self._chunk):
stop = min(start + self._chunk, X.shape[0])
scores = X[start:stop] @ centroids.T
scores *= -2.0
scores += centroid_norms
out[start:stop] = scores.argmin(axis=1)
return outWhat changed: The broadcast assignment materializes an (n, k, d) intermediate, which at a million points and a hundred clusters is tens of gigabytes — the memory, not the arithmetic, is what breaks. Expanding the squared distance removes it: only the cross term touches both operands, so the assignment becomes a GEMM, and the per-point norm is constant within a row and is never computed because it cannot change the argmin. Chunking bounds the remaining score block, and the M-step is rewritten as sufficient statistics accumulated across chunks, so the full assignment never exists at once.
Vectorize to NumPy/BLAS (@, np.dot, einsum)
The cross term becomes one GEMM per chunk and the centroid norms an einsum reduction computed once per iteration rather than per point.
Trade-off: The expansion is less numerically stable than a direct subtraction — cancellation between large norms can misrank two nearly-equidistant centroids, which matters when clusters overlap.
Batch work to amortize interpreter overhead
Chunking bounds the score block at (chunk, k) instead of (n, k), so the fit runs on datasets far larger than that matrix while still entering BLAS only once per chunk.
Trade-off: The M-step must be restated as accumulated sufficient statistics, which is less obvious than a single grouped mean, and the chunk size is a tuning parameter with no good default.
Pre-allocate output arrays and use in-place operations
Totals, counts and labels are allocated once for the whole fit, and the score block is built with in-place multiply and add rather than three temporaries per chunk.
Trade-off: The in-place chain is order-dependent and the buffers are shared across iterations, so the object is not reentrant and a reordered line corrupts the assignment silently.
O(n*k*d) per iteration in BLAS, memory bounded by the chunk rather than by n. Illustrative, not a measured benchmark.