Skip to content
DeepForgeby svx
All problems
MediumML Fundamentalsml-032

K-Means Assignment Step

Problem

Assignment step of k-means: assign each point to the index of its nearest centroid by Euclidean distance (comparing squared distances is equivalent).

Return one cluster index per point; ties go to the smallest centroid index.

Hint

Use strict less-than when updating the best centroid so ties keep the smallest index.

Test cases

case 1
in: [[[0,0],[0,1],[10,10],[10,11]],[[0,0],[10,10]]]
expected: [0,0,1,1]
case 2
in: [[[1,1],[8,8],[9,10]],[[0,0],[5,5],[10,10]]]
expected: [0,2,2]
case 3
in: [[[3,3]],[[0,0],[10,10]]]
expected: [0]
case 4
in: [[[1,1],[2,2]],[[2,2],[1,1]]]
expected: [1,0]
Starter codePython
def kmeans_assign(points, centroids):
    # Your code here
    pass

4 test cases · run your code against them in the editor