K-Means: centroid update

The problem, from first principles

01-kmeans-assignment answered "given wherever the centroids currently sit, which one is each point closest to." That answer is only ever as good as the guessed centroids it started from — if the centroids are in the wrong place, the assignment is wrong too. The fix is the obvious one: now that you know which points think they belong to which centroid, move each centroid to where those points actually are.

That "move to where your assigned points actually are" step is the same centroid computation 04-nearest-centroid did for labeled classes — average the points in a group to get the group's center — just applied to a group defined by "current assignment" instead of "true class." Alternate this with the assignment step and the two together are the whole K-Means algorithm.

From theory to code

Implement kmeans_update_centroids(input, assignments, centroids): for each cluster index, average the rows of input currently assigned to it, and use that average as the cluster's new centroid. The one piece of theory that isn't just "average a group" is what to do when a group is empty — that's called out separately in Constraints and Theory below.

Constraints

• input: shape (n_samples, n_features).
• assignments: shape (n_samples,), integer cluster index per sample, as produced by 01-kmeans-assignment's kmeans_assign.
• centroids: shape (k, n_features), the current centroids.
• Output: shape (k, n_features).
• A cluster with zero points assigned to it keeps its old centroid unchanged — never average an empty group (0/0 is undefined).
• The input centroids array must not be mutated in place; callers may still hold a reference to it.

Hints

Open one at a time. Each gives away a little more than the last.