Neural Mastery
← Practice
K-Means Clustering: Single Lloyd Stepmedium

K-Means Clustering: Single Lloyd Step

medium Classical ML from Scratch 15–20 min

NumPy allowed · plain Python earns 10 bonus points

Execute a single iteration of Lloyd's algorithm for K-means clustering.

Related Skills

Task

Implement kmeans_step(X, centroids): assign each point to its nearest centroid by squared Euclidean distance, then update each centroid to the mean of assigned points.

Function Signature

kmeans_step(X: list[list[float]], centroids: list[list[float]]) -> tuple[list[int], list[list[float]]]

Examples

Example 1: 2 Clusters in 2D
Input: X = [[1, 1], [1.5, 2], [5, 5], [6, 5]] centroids = [[0, 0], [10, 10]]
Output: [[0, 0, 0, 1], [[2.5, 2.6667], [6, 5]]]
Example 2: k=1 Single Cluster
Input: X = [[2, 4], [4, 6]] centroids = [[0, 0]]
Output: [[0, 0], [[3, 5]]]

Constraints

  • 1 <= len(X) <= 1000
  • 1 <= len(centroids) <= 50
  • len(X[0]) == len(centroids[0])
Python3Saved ✓
def kmeans_step(X: list[list[float]], centroids: list[list[float]]) -> tuple[list[int], list[list[float]]]:
"""X: data points (n x d).
centroids: current cluster centers (k x d).
Perform one step of Lloyd's algorithm:
1. Assign each point to closest centroid (break ties by smallest centroid index).
2. Recompute each centroid as mean of its assigned points (rounded to 4 decimal places).
(If a centroid has no points, keep it unchanged).
Return (assignments, new_centroids)."""
# Your implementation here
pass

Case 1: 2 Clusters in 2D
Input: X = [[1, 1], [1.5, 2], [5, 5], [6, 5]] centroids = [[0, 0], [10, 10]]
Expected: [[0, 0, 0, 1], [[2.5, 2.6667], [6, 5]]]
Case 2: k=1 Single Cluster
Input: X = [[2, 4], [4, 6]] centroids = [[0, 0]]
Expected: [[0, 0], [[3, 5]]]