K-Means: One Assignment Step

Difficulty: Medium · Concept: K-Means & Hierarchical Clustering

K-Means alternates two steps until convergence: assign each point to its nearest centroid, then update each centroid to the mean of its assigned points. This problem is the assign step in isolation — the piece that decides cluster membership by plain Euclidean distance.

Your task: implement kmeans_assignment_step(points, centroids). points and centroids are lists of equal-length numeric coordinate lists (works for any number of dimensions, not just 2D). Return a list of integers — the index of the nearest centroid for each point, in the same order as points. On a tie, assign to the lowest-index centroid.

Implement it yourself

assert kmeans_assignment_step([[0, 0], [0, 1], [5, 5], [5, 6]], [[0, 0], [5, 5]]) == [0, 0, 1, 1]
assert kmeans_assignment_step([[2, 2]], [[0, 0], [10, 10]]) == [0]
assert kmeans_assignment_step([[1, 0]], [[0, 0], [2, 0]]) == [0]
assert kmeans_assignment_step([[1], [9]], [[0], [10]]) == [0, 1]

Next: K-Means: Centroid Update Step (a harder variant), or skip ahead to Sigmoid Activation From Scratch