Implement K-means with Lloyd's algorithm. What does it optimise, how do you initialise it, and what do you do when inference data doesn't look like training data?

import numpy as np

def kmeans(X, k, iters=100, seed=0):
    rng = np.random.default_rng(seed)
    C = X[rng.choice(len(X), k, replace=False)]
    for _ in range(iters):
        d = ((X[:, None] - C[None]) ** 2).sum(-1)
        a = d.argmin(1)
        new = np.array([X[a == j].mean(0)
                        if (a == j).any() else C[j]
                        for j in range(k)])
        if np.allclose(new, C):
            break
        C = new
    return C, a

It minimises the within-cluster sum of squared distances: assign points to the nearest centroid, move centroids to their means, repeat. It reaches a local optimum, so use k-means++ and restarts. For train-inference mismatch, reuse the training scaler and monitor distances to centroids; drift means refit.