K-Means: Assign to Nearest Centroid - Unsupervised Learning

K-Means: Assign to Nearest Centroid

What you'll build

The assignment step of k-means: given data points and centroids, return the index of the nearest centroid for every point, without a Python loop over points.

Background

k-means alternates two steps: assign each point to its nearest centroid, then update each centroid to the mean of its points. The assignment step is where almost all of the compute goes, O(n * k * d) per iteration, so vectorising it is the main skill being tested. The same pairwise-distance pattern powers k-NN and vector search. (The full algorithm is K-Means from Scratch later in this section.)

Formula

d(i, j)  = sqrt( sum_f (X[i, f] - centroids[j, f])^2 )      # Euclidean distance
label[i] = argmin_j d(i, j)

The square root doesn't change which centroid is nearest, so comparing squared distances gives the same labels.

Inputs

• X: data points, shape (n_samples, n_features)
• centroids: cluster centres, shape (k, n_features)

(Tests may pass nested Python lists; np.asarray them.)

Expected return format

An integer array of shape (n_samples,) with values in 0 .. k-1.

Edge cases and conventions

• Ties: a point equally far from several centroids goes to the lowest centroid index (what np.argmin returns).

Constraints

• NumPy only, vectorised with broadcasting: X[:, None, :] - centroids[None, :, :] has shape (n, k, d).

Follow-up

• For large n and d, avoid the (n, k, d) intermediate with ||x||^2 - 2 x.c + ||c||^2 (one matrix multiply). What numerical issue can that introduce?

Input: assign_clusters(np.array([[0, 0], [1, 1], [10, 10]]), np.array([[0, 0], [10, 10]]))
Output: [0, 0, 1]

Input: assign_clusters(np.array([[5, 5]]), np.array([[0, 0], [10, 10]]))
Output: [0]

import numpy as np

def assign_clusters(X: np.ndarray, centroids: np.ndarray) -> np.ndarray:
    """
    Assign each point to its nearest centroid (Euclidean distance).

    Args:
        X: Data points, shape (n_samples, n_features)
        centroids: Cluster centres, shape (k, n_features)

    Returns:
        Int array of shape (n_samples,) with values in 0..k-1
        (ties go to the lowest centroid index)
    """
    # Your code here
    pass