k-Means Clustering From Scratch in NumPy

Run k-means on X (n × d) from the starting centroids init (k × d):

Assign every point to its nearest centroid (squared Euclidean distance; ties go to the lower centroid index).
Repeat up to max_iter times: move each centroid to the mean of its assigned points (a centroid with no points stays where it is); reassign every point; stop early if no assignment changed.

Return [centroids, labels], where the labels are the assignment to the returned centroids.

Examples

Input X = [[0, 0], [0, 1], [10, 10], [10, 11]], init = [[0, 0], [10, 10]], max_iter = 10

Output [[[0, 0.5], [10, 10.5]], [0, 0, 1, 1]]

Two clean clusters. One update and the centroids sit at the cluster means.

Input X = [[0, 0], [1, 0], [2, 0]], init = [[0, 0], [100, 100]], max_iter = 10

Output [[[1, 0], [100, 100]], [0, 0, 0]]

The second centroid never wins a point, so it stays put.

Input X = [[0, 0], [4, 0]], init = [[1, 0], [3, 0]], max_iter = 0

Output [[[1, 0], [3, 0]], [0, 1]]

max_iter = 0: return the starting centroids and the first assignment.

Constraints

1 ≤ n ≤ 5000, 1 ≤ d ≤ 20, 1 ≤ k ≤ 50
0 ≤ max_iter ≤ 300
Integer coordinates
Checked to a relative tolerance of 10⁻⁹

Hints

Hint 1

Assignment is an argmin over an (n, k) matrix of squared distances. np.argmin returns the first index on ties.

Hint 2

For the update, X[labels == j].mean(axis=0) for each j, skipping clusters with no members.

Hint 3

After each update, reassign; if the new labels equal the old ones, stop. Otherwise keep the new labels and continue.

Where candidates go wrong

Calling .mean() on an empty selection: NumPy returns nan with a warning and the centroid is lost.
Checking convergence on the centroids with ==, which can fail to trigger because of rounding. Compare the assignments.
Returning labels from the previous assignment after the final update, so they do not match the centroids.

The idea

Each iteration is two vectorised steps. Assignment: build the (n, k) squared distance matrix by broadcasting and take argmin along axis 1. Update: for each centroid, average the points assigned to it. Both steps can only reduce the within-cluster sum of squares, and there are finitely many assignments, so the loop terminates; comparing assignments is the exact way to detect it.

An empty cluster has no mean, so the rule here is to leave its centroid unchanged. scikit-learn instead moves it to the point farthest from its current centre.

Time O(iterations · nkd) · Space O(nk) for the distance matrix

Reference solution

import numpy as np

def kmeans(X, init, max_iter):
    X = np.asarray(X, dtype=float)
    C = np.array(init, dtype=float)

    def assign(C):
        D = ((X[:, None, :] - C[None, :, :]) ** 2).sum(axis=2)   # (n, k)
        return D.argmin(axis=1)                                  # first index on ties

    labels = assign(C)
    for _ in range(max_iter):
        for j in range(len(C)):
            members = X[labels == j]
            if len(members):                                     # an empty cluster keeps its centroid
                C[j] = members.mean(axis=0)
        new = assign(C)
        if np.array_equal(new, labels):
            break
        labels = new
    return [C, labels]

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

The standard unsupervised implementation question. A good answer is vectorised, stops at the right moment, and has decided in advance what to do with an empty cluster. The usual follow-up is k-means++.

The theory behind it is in the lesson k-means, kNN, PCA and decision trees from scratch.

Keep going

Easy12 min

k-means clustering (Lloyd's algorithm)

Examples

Constraints

Hints

Where candidates go wrong

The idea

Reference solution

Follow-ups an interviewer asks

Why interviewers ask it

Keep going

k-nearest neighbours classifier

k-means clustering (Lloyd's algorithm)

Principal component analysis

Best decision-tree split by Gini impurity