Initialize by k-means++ (probability proportional to ||x − nearest centroid||²), then iterate Lloyd's algorithm.
import numpy as np
def kmeans(X, k, max_iter=100, tol=1e-6, seed=0):
rng = np.random.default_rng(seed)
centroids = X[rng.choice(len(X), k, replace=False)] # or kmeans++ init
for _ in range(max_iter):
dists = ((X[:, None, :] - centroids[None, :, :]) ** 2).sum(-1)
labels = dists.argmin(axis=1)
new_centroids = np.array([
X[labels == j].mean(axis=0) if (labels == j).any()
else X[rng.integers(len(X))] # re-seed empty cluster
for j in range(k)
])
if np.linalg.norm(new_centroids - centroids) < tol:
break
centroids = new_centroids
return centroids, labels
Gotchas: empty cluster after an update (re-seed from the farthest point); label drift from lazily using assignments inside mean(); broadcasting mistakes in ||X − c||²; float32 overflow for >1M points. Mention k-means++ without being asked — it's the expected "senior" flourish. Direct AI tie-in: k-means is exactly the coarse quantizer in IVF vector indexes — retrieval companies love this connection.