K-Means Clustering

Implement Lloyd's algorithm, the standard k-means loop.

Starting from initial_centroids of shape (k, d), repeat iters times:

1. Assign. Give every point the index of the nearest centroid by Euclidean distance. Break ties towards the lower centroid index.
2. Update. Move each centroid to the mean of the points assigned to it. If a cluster ends up empty, leave that centroid exactly where it was.

Return (centroids, labels) where centroids has shape (k, d) and labels has shape (n,). The labels must be the assignment computed against the returned centroids, so run one final assignment step after the last update.

Real k-means is randomised through its initialisation. Taking the centroids as an argument makes the result deterministic and testable, and it is also how k-means++ plugs in.

Function signature

def kmeans(X, initial_centroids, iters):

Example

Input

X = np.array([[0.0], [0.5], [10.0], [10.5]])
kmeans(X, np.array([[0.0], [1.0]]), iters=5)

Output

(array([[ 0.25], [10.25]]), array([0, 0, 1, 1]))

Reasoning

In the first pass, 0.5 is equally far from both centroids, so it goes to the lower index, 0. The large points go to centroid 1. The centroids move to 0.25 and 10.25, and after that nothing changes.

Constraints and notes

• Break distance ties towards the lower centroid index.
• An empty cluster keeps its centroid where it was.
• Assign once more after the last update.