K-Means Clustering

Implement k-means clustering (Lloyd’s algorithm) from scratch on a 2D array of points. K-means partitions n points into k clusters by alternating two steps: assign each point to its nearest centroid, then move each centroid to the mean of the points assigned to it. It is a staple of unsupervised-learning interviews, and the interviewer is checking that you can vectorize the assignment step instead of writing a distance loop.

Algorithm

1. Initialize k centroids by sampling k distinct rows of X at random.
2. Assign step: label each point with the index of its nearest centroid (Euclidean distance).
3. Update step: set each centroid to the mean of the points currently assigned to it.
4. Repeat steps 2-3 until the labels stop changing or you reach max_iters.

Task

Complete kmeans(X, k, max_iters=100, seed=0) so it runs Lloyd’s algorithm and returns (centroids, labels).

• Draw the initial centroids by choosing k distinct row indices of X using a local np.random.RandomState(seed) (for example rng.choice(n, size=k, replace=False)). Use only this local RNG so results are reproducible.
• In each iteration, compute the (n, k) matrix of Euclidean distances from every point to every centroid with vectorized NumPy (no Python loop over points), then take argmin along the centroid axis to get the new labels.
• Recompute each centroid as the mean of its assigned points. If a cluster ends up empty, keep its previous centroid so the mean is well defined.
• Stop early when the new labels equal the previous labels; otherwise stop after max_iters iterations. Return the final centroids (shape (k, d), float) and labels (shape (n,), int).

Do not call sklearn, scipy, or any built-in k-means / pairwise-distance helper.

Example

X = np.array([[0.0, 0.0],
              [0.1, 0.0],
              [0.0, 0.1],
              [5.0, 5.0],
              [5.1, 5.0],
              [5.0, 5.1]])
centroids, labels = kmeans(X, k=2, max_iters=100, seed=0)
# The two dense blobs are recovered; one label per point, e.g.
# labels -> array([0, 0, 0, 1, 1, 1])   (cluster ids may be swapped)
# centroids ~ [[0.033, 0.033], [5.033, 5.033]]  (order matches labels)

The exact cluster ids (0 vs 1) depend on which rows the seed picks first, but the partition into the two blobs and the centroid means are determined.

Constraints

• 1 <= k <= n, X has shape (n, d) with 1 <= n <= 10^4, 1 <= d <= 10^3.
• Use vectorized NumPy for the assignment step; no Python loop over the n points.
• Seed all randomness through a local np.random.RandomState(seed); do not modify the global NumPy RNG.
• Handle empty clusters without crashing (a naive mean of an empty slice returns nan).

Approach

Run Lloyd’s algorithm: pick k distinct rows of X as the initial centroids using a seeded local RNG, then alternate assign and update. The assign step is a single vectorized pairwise-distance computation (squared-norm expansion) followed by argmin over the k centroids, so there is no Python loop over points. The update step recomputes each centroid as the mean of its members, keeping the old centroid for any empty cluster; iteration stops as soon as the labels are unchanged or max_iters is hit.

Solution

import numpy as np

def kmeans(X: np.ndarray, k: int, max_iters: int = 100, seed: int = 0):
    X = np.asarray(X, dtype=np.float64)
    n = X.shape[0]

    # 1. Initialize centroids from k distinct points using a local RNG.
    rng = np.random.RandomState(seed)
    init_idx = rng.choice(n, size=k, replace=False)
    centroids = X[init_idx].copy()

    labels = np.full(n, -1, dtype=int)

    for _ in range(max_iters):
        # 2. Assign: (n, k) squared distances, then nearest centroid.
        #    ||x - c||^2 = ||x||^2 - 2 x.c + ||c||^2; the ||x||^2 term is
        #    constant across centroids and does not change the argmin.
        cross = X @ centroids.T                       # (n, k)
        c_sq = np.sum(centroids ** 2, axis=1)         # (k,)
        sq_dist = -2.0 * cross + c_sq[None, :]        # (n, k), up to +||x||^2
        new_labels = np.argmin(sq_dist, axis=1)

        # 4. Converged when no point changed cluster.
        if np.array_equal(new_labels, labels):
            labels = new_labels
            break
        labels = new_labels

        # 3. Update: each centroid becomes the mean of its members;
        #    empty clusters keep their previous centroid.
        for j in range(k):
            members = X[labels == j]
            if members.shape[0] > 0:
                centroids[j] = members.mean(axis=0)

    return centroids, labels

Walkthrough

On the example with the two blobs near (0, 0) and (5, 5) and seed=0:

• rng.choice(6, size=2, replace=False) picks two distinct row indices; say one lands in each blob (if both land in the same blob the algorithm still converges in a couple more iterations).
• Assign step: X @ centroids.T gives each point’s dot product with both centroids; adding c_sq and taking argmin labels every point with its nearer centroid. Points 0-2 fall to the (0,0) centroid, points 3-5 to the (5,5) centroid.
• Update step: centroid 0 becomes the mean of rows 0-2 = [0.033, 0.033], centroid 1 the mean of rows 3-5 = [5.033, 5.033].
• Next iteration produces the same labels, so np.array_equal is true and the loop breaks, returning those centroids and labels = [0,0,0,1,1,1] (ids depend on which blob the first sampled row was in).

Complexity & notes

• Time O(iters * n * k * d), dominated by the X @ centroids.T matmul per iteration (BLAS in C). Space O(n*k) for the distance matrix plus O(k*d) for the centroids.
• Dropping the ||x||^2 term is a deliberate optimization: it is the same for every centroid of a given point, so it never affects the per-row argmin. Add it back (and take sqrt) only if you need the actual distances, not just the assignment.
• The empty-cluster guard matters: X[labels == j].mean(axis=0) over an empty slice returns nan, which then poisons every future distance. Keeping the previous centroid is the simplest fix; production code often reseeds an empty centroid to the point farthest from any center.
• Seeding through a local np.random.RandomState(seed) keeps runs reproducible and avoids mutating the global NumPy RNG that the rest of a program may depend on.
• The k loop in the update step runs over clusters, not points, so it is O(k) iterations of vectorized means and does not violate the no-loop-over-points requirement. Lloyd’s algorithm finds a local optimum only; real use runs several seeds and keeps the lowest inertia (k-means++ improves initialization).