Implement k-means clustering from scratch

Implement k-means clustering from scratch. Given a data matrix X of n points with d features each and a number of clusters k, return the final centroids and the cluster label of every point. Make each step unambiguous and keep the matrix shapes consistent.

Interface

Unless you agree on a different interface with the interviewer, assume:

• kmeans(X: list[list[float]], k: int, max_iters: int = 100) -> tuple[list[list[float]], list[int]] — returns (centroids, labels).

To make results deterministic:

• Start from the first k points of X as the centroids.
• Assign each point to the nearest centroid by squared Euclidean distance; on a tie, the smaller centroid index wins.
• Move each centroid to the mean of its points. A cluster with no points keeps its centroid.
• Stop when no label changes, or after max_iters iterations.

Example

kmeans([[1, 1], [1, 2], [4, 4], [4, 5]], 2)
  -> ([[1.0, 1.5], [4.0, 4.5]], [0, 0, 1, 1])
# start (1,1), (1,2): first pass labels [0, 1, 1, 1], centroids (1,1) and (3, 11/3);
# second pass labels [0, 0, 1, 1]; third pass changes nothing

kmeans([[0, 0], [2, 2], [4, 4]], 1)
  -> ([[2.0, 2.0]], [0, 0, 0])

Constraints

• 1 <= n <= 1000, 1 <= d <= 20, 1 <= k <= n, 1 <= max_iters <= 100.

Clarifying questions

• Plain Python, or NumPy?
• How should centroids be initialised: given, random with a seed, or k-means++?
• What should happen to a cluster that loses all its points?

What a strong answer covers

• The two alternating steps (assignment and update) and why each one can only lower the within-cluster sum of squares, so the loop converges, though only to a local optimum.
• Shapes in a vectorised version: distances (n, k) from X (n, d) and C (k, d) with ||x||² - 2x·c + ||c||², labels by argmin, centroids from per-cluster sums and counts.
• Cost per iteration O(n · k · d).
• Initialisation (k-means++, several restarts), empty clusters, choosing k (elbow, silhouette), and feature scaling.

Follow-ups

• Implement k-means++ initialisation.
• The data does not fit in memory. Use mini-batch k-means.
• When does k-means do badly, and what would you use instead?