K-Means Centroid Update

Problem

Centroid update step of k-means: average the points assigned to each of k clusters.

If a cluster has no assigned points, keep old_centroids[i] unchanged. Return the list of k new centroids.

Hint

Group by assignment, average each group coordinate-wise, and copy old centroids for empty clusters.

case 1

in: [[[0,0],[0,1],[10,10],[10,11]],[0,0,1,1],2,[[0,0],[10,10]]]

expected: [[0,0.5],[10,10.5]]

case 2

in: [[[1,1],[8,8],[9,10]],[0,1,1],2,[[0,0],[5,5]]]

expected: [[1,1],[8.5,9]]

case 3

in: [[[2,2],[4,4]],[0,0],3,[[0,0],[5,5],[9,9]]]

expected: [[3,3],[5,5],[9,9]]

case 4

in: [[[1],[3],[10]],[0,1,1],2,[[0],[0]]]

expected: [[1],[6.5]]

Starter codePython

def kmeans_update_centroids(points, assignments, k, old_centroids):
    # Your code here
    pass

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