K-Means: Centroid Update Step
Difficulty: Medium · Concept: K-Means & Hierarchical Clustering
K-Means: One Assignment Step implemented the assign half of Lloyd's algorithm — nearest-centroid lookup. This is the other half: update — given the current assignments, recompute each centroid as the mean of the points now assigned to it. Together, these two steps are the entire K-Means loop.
The genuine edge case a plain "average the points" implementation misses: a cluster can end up with zero points assigned to it (a bad random initialization, or a centroid that got pushed somewhere no point is closer to). Averaging zero points is a division by zero — a real implementation has to decide what happens then, and the standard, sane choice is: leave that centroid exactly where it was, rather than crash or silently produce nan.
Your task: implement kmeans_update_step(points, assignments, k, old_centroids). assignments[i] is the cluster index (0 to k-1) for points[i]. Return a list of k new centroid positions.