Assign Points to Nearest Centroid

Easy
~15 min
code completion

K-Means: Assignment Step

K-means alternates between two steps. The first is assignment: assign each point to its nearest centroid.

For each data point , find the centroid that minimizes the Euclidean distance:

In NumPy, compute all distances at once using broadcasting:

# X: (n, d),  centroids: (k, d)
dists = np.linalg.norm(X[:, None] - centroids[None, :], axis=2)  # (n, k)
labels = np.argmin(dists, axis=1)  # (n,)

Your task:

Implement assign_clusters(X, centroids) that returns the cluster label for each point.

Example Tests

Two clear clusters

Input: {"X":[[0,0],[0.1,0.1],[10,10],[10.1,9.9]],"centroids":[[0,0],[10,10]]}

Expected: [0,0,1,1]

Point equidistant picks lower index (argmin)

Input: {"X":[[5,0]],"centroids":[[0,0],[10,0]]}

Expected: [0]

3 centroids

Input: {"X":[[1,0],[5,0],[9,0]],"centroids":[[0,0],[5,0],[10,0]]}

Expected: [0,1,2]

Sign in to solve this problem

You can read the full problem statement above. Create a free account to run code in the browser, submit solutions, and track your progress.