Implement k-means from scratch, including k-means++ initialization and a convergence check.

TL;DR: k-means alternates two steps until convergence: assign each point to its nearest centroid, then move each centroid to the mean of its assigned points. Initialize with k-means++ (probability proportional to squared distance from existing centers) so you avoid bad local minima, and stop when centroids barely move or after a max-iteration cap. It minimizes within-cluster squared distance and only finds a local optimum.

How to approach it. State the objective (minimize within-cluster sum of squares), name the two alternating steps, and call out the two things juniors skip: principled initialization (k-means++) and a real stopping rule. Then write it vectorized.

A strong answer. Lloyd's algorithm with k-means++ init:

import numpy as np

def kmeans(X, k, max_iter=100, tol=1e-4, seed=0):
    rng = np.random.default_rng(seed)
    # --- k-means++ init: spread initial centers out ---
    centers = [X[rng.integers(len(X))]]
    for _ in range(1, k):
        d2 = np.min(((X[:, None] - np.array(centers)) ** 2).sum(-1), axis=1)
        probs = d2 / d2.sum()                      # farther points more likely
        centers.append(X[rng.choice(len(X), p=probs)])
    centers = np.array(centers)

    for _ in range(max_iter):
        d2 = ((X[:, None] - centers) ** 2).sum(-1)  # (n, k) squared distances
        labels = d2.argmin(1)
        new = np.array([X[labels == j].mean(0) if np.any(labels == j)
                        else centers[j] for j in range(k)])  # keep empty clusters put
        if np.linalg.norm(new - centers) < tol:     # converged
            centers = new
            break
        centers = new
    return labels, centers

Three details that signal experience: k-means++ picks each new center with probability proportional to squared distance from the nearest existing center, which sharply cuts the chance of a bad local optimum versus random init; the convergence check on centroid movement (not a fixed iteration count alone); and empty-cluster handling so a centroid with no points does not produce a NaN. The algorithm always converges but only to a local minimum, so in practice you run it several times with different seeds and keep the lowest inertia.

Key takeaways

• k-means++ spreads initial centers by squared-distance sampling, avoiding the bad local minima random init falls into.
• Stop on centroid movement below tol, not iteration count alone, and keep a max-iter safety cap.
• Handle empty clusters explicitly or a mean over zero points returns NaN.
• Standardize features first: unscaled dimensions dominate the Euclidean distance and skew every cluster.

What interviewers probe next.

• "How do you choose k?" Elbow on inertia or the silhouette score; in production, tie it to the downstream use, not just the curve.
• "Complexity?" O(n·k·d) per iteration. For large n, mini-batch k-means or a KD-tree/approximate assignment.
• "When does k-means fail?" Non-convex or differently-sized/density clusters, and unscaled features (it assumes Euclidean geometry, so standardize first). Use DBSCAN or GMM when clusters are not spherical.
• "Why squared distance, not absolute?" The mean is the minimizer of squared distance; with L1 you would use the median (k-medians).

Common mistakes.

• Random initialization with no restarts, landing in a bad local optimum.
• No empty-cluster handling, producing NaN centroids.
• Stopping on iteration count alone with no movement-based convergence test.
• Forgetting to standardize features, so one large-scale feature dominates the distance.