Easy

Update Centroids

Easy

~15 min

code completion

K-Means: Update Step

After assigning all points to clusters, the second step is centroid update: move each centroid to the mean of its assigned points.

μ_{k} = \frac{1}{∣ C _{k} ∣} i \in C_{k} \sum x_{i}

In NumPy, for each cluster $k$ , select all rows of $X$ where labels == k and take their mean:

new_centroids[k] = X[labels == k].mean(axis=0)

Your task:

Implement update_centroids(X, labels, k) that returns the new centroid matrix of shape (k, d).

Example Tests

Two clusters: centroids at their means

Input: {"X":[[0,0],[2,0],[10,0],[12,0]],"k":2,"labels":[0,0,1,1]}

Expected: [[1,0],[11,0]]

All points in one cluster

Input: {"X":[[1,1],[3,3],[5,5]],"k":1,"labels":[0,0,0]}

Expected: [[3,3]]

Output shape is (k, d)

Input: {"X":[[1,2],[3,4],[5,6]],"k":3,"labels":[0,1,2]}

Expected: [3,2]

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