Lloyd's k-means reads the whole dataset before it moves a single centroid. e-means, a 2026 method from Meta FAIR, moves the centroids after every mini-batch instead, and it reaches the loss that Lloyd's algorithm needs 25 epochs for in about two. This problem asks you to implement its update rule. The two k-means problems in the catalog cover the assignment step and the centroid update, and this problem combines them inside a mini-batch loop.
Every cluster j keeps two running statistics, a sum S_j of shape (d,) and a count N_j, and its centroid is
cj=NjSjAt the start, each cluster behaves as if it held exactly one point at its initial centroid, so S_j is the initial centroid and N_j is 1.
The batches are processed in the order given. For each batch, assign every point in it to its nearest centroid by squared Euclidean distance. Let s_j be the sum of the batch points assigned to cluster j and n_j their number. Both statistics then move toward the batch values with step alpha,
Sj←(1−α)Sj+αsj,Nj←(1−α)Nj+αnjand every centroid becomes S_j divided by N_j before the next batch is assigned. A cluster that receives no points in a batch has both statistics shrunk by the same factor, so its centroid stays where it was while its count fades.
After the last batch, assign every point of X to its nearest final centroid. Return a dict with four entries. centroids is a float NumPy array of shape (k, d) with the final centroids. counts is a float NumPy array of shape (k,) with the final running counts N. labels is an integer NumPy array of shape (n,) with the final assignment of every point. mse is a Python float, the mean over all n points of the squared distance to their assigned final centroid.
Worked example. Take the points (0, 0), (1, 0), (9, 0) and (10, 0), the starting centroids (2, 0) and (8, 0), alpha equal to 0.5 and the batches [0, 2] and [1, 3]. In the first batch (0, 0) joins cluster 0 and (9, 0) joins cluster 1. Cluster 0 gets S equal to half of (2, 0) plus half of (0, 0), which is (1, 0), and N equal to half of 1 plus half of 1, which is 1, so its centroid moves to (1, 0). Cluster 1 moves to half of 8 plus half of 9 on the first axis, which is 8.5. In the second batch (1, 0) joins cluster 0, whose centroid stays at (1, 0), and (10, 0) pulls cluster 1 to 9.25. The final labels are [0, 0, 1, 1] and the mse is (1 + 0 + 0.0625 + 0.5625) divided by 4, which is 0.40625.
Input: X = [[0, 0], [1, 0], [9, 0], [10, 0]], centroids = [[2, 0], [8, 0]], batches = [[0, 2], [1, 3]], alpha = 0.5
Output: centroids = [[1.0, 0.0], [9.25, 0.0]], counts = [1.0, 1.0], labels = [0, 0, 1, 1], mse = 0.40625
Explanation: Each batch gives one point to each cluster, so the counts stay at 1 while each centroid moves halfway toward the new point.
Input: X = [[0, 0], [1, 0], [9, 0], [10, 0]], centroids = [[2, 0], [8, 0]], batches = [[0, 1, 2, 3]], alpha = 1.0
Output: centroids = [[0.5, 0.0], [9.5, 0.0]], counts = [2.0, 2.0], labels = [0, 0, 1, 1], mse = 0.25
Explanation: With one batch holding every point and a step of 1, the old statistics are forgotten completely and each centroid is the exact mean of its points, which is one step of Lloyd's algorithm.
Input: X = [[0, 0], [2, 0], [20, 0]], centroids = [[3, 0], [19, 0]], batches = [[0, 1]], alpha = 0.5
Output: centroids = [[1.666667, 0.0], [19.0, 0.0]], counts = [1.5, 0.5], labels = [0, 0, 1], mse = 1.296296
Explanation: Both batch points join cluster 0, so its sum becomes (1.5, 0) plus (1, 0) and its count 0.5 plus 1. Cluster 1 gets no points, so its sum and count both halve and its centroid stays at (19, 0).
((xb[:, None, :] - C[None, :, :]) ** 2).sum(axis=-1) gives a (B, k) distance table, and np.argmin(dist, axis=1) returns the first smallest index.
np.add.at(s, a, xb) sums the batch points into their clusters, and np.bincount(a, minlength=k) counts them.
C[N > 0] = S[N > 0] / N[N > 0][:, None] updates only the clusters with a nonzero count.
Accepts: array
Accepts: array
Accepts: array
Accepts: number
Lloyd's k-means reads the whole dataset before it moves a single centroid. e-means, a 2026 method from Meta FAIR, moves the centroids after every mini-batch instead, and it reaches the loss that Lloyd's algorithm needs 25 epochs for in about two. This problem asks you to implement its update rule. The two k-means problems in the catalog cover the assignment step and the centroid update, and this problem combines them inside a mini-batch loop.
Every cluster j keeps two running statistics, a sum S_j of shape (d,) and a count N_j, and its centroid is
cj=NjSjAt the start, each cluster behaves as if it held exactly one point at its initial centroid, so S_j is the initial centroid and N_j is 1.
The batches are processed in the order given. For each batch, assign every point in it to its nearest centroid by squared Euclidean distance. Let s_j be the sum of the batch points assigned to cluster j and n_j their number. Both statistics then move toward the batch values with step alpha,
Sj←(1−α)Sj+αsj,Nj←(1−α)Nj+αnjand every centroid becomes S_j divided by N_j before the next batch is assigned. A cluster that receives no points in a batch has both statistics shrunk by the same factor, so its centroid stays where it was while its count fades.
After the last batch, assign every point of X to its nearest final centroid. Return a dict with four entries. centroids is a float NumPy array of shape (k, d) with the final centroids. counts is a float NumPy array of shape (k,) with the final running counts N. labels is an integer NumPy array of shape (n,) with the final assignment of every point. mse is a Python float, the mean over all n points of the squared distance to their assigned final centroid.
Worked example. Take the points (0, 0), (1, 0), (9, 0) and (10, 0), the starting centroids (2, 0) and (8, 0), alpha equal to 0.5 and the batches [0, 2] and [1, 3]. In the first batch (0, 0) joins cluster 0 and (9, 0) joins cluster 1. Cluster 0 gets S equal to half of (2, 0) plus half of (0, 0), which is (1, 0), and N equal to half of 1 plus half of 1, which is 1, so its centroid moves to (1, 0). Cluster 1 moves to half of 8 plus half of 9 on the first axis, which is 8.5. In the second batch (1, 0) joins cluster 0, whose centroid stays at (1, 0), and (10, 0) pulls cluster 1 to 9.25. The final labels are [0, 0, 1, 1] and the mse is (1 + 0 + 0.0625 + 0.5625) divided by 4, which is 0.40625.
Input: X = [[0, 0], [1, 0], [9, 0], [10, 0]], centroids = [[2, 0], [8, 0]], batches = [[0, 2], [1, 3]], alpha = 0.5
Output: centroids = [[1.0, 0.0], [9.25, 0.0]], counts = [1.0, 1.0], labels = [0, 0, 1, 1], mse = 0.40625
Explanation: Each batch gives one point to each cluster, so the counts stay at 1 while each centroid moves halfway toward the new point.
Input: X = [[0, 0], [1, 0], [9, 0], [10, 0]], centroids = [[2, 0], [8, 0]], batches = [[0, 1, 2, 3]], alpha = 1.0
Output: centroids = [[0.5, 0.0], [9.5, 0.0]], counts = [2.0, 2.0], labels = [0, 0, 1, 1], mse = 0.25
Explanation: With one batch holding every point and a step of 1, the old statistics are forgotten completely and each centroid is the exact mean of its points, which is one step of Lloyd's algorithm.
Input: X = [[0, 0], [2, 0], [20, 0]], centroids = [[3, 0], [19, 0]], batches = [[0, 1]], alpha = 0.5
Output: centroids = [[1.666667, 0.0], [19.0, 0.0]], counts = [1.5, 0.5], labels = [0, 0, 1], mse = 1.296296
Explanation: Both batch points join cluster 0, so its sum becomes (1.5, 0) plus (1, 0) and its count 0.5 plus 1. Cluster 1 gets no points, so its sum and count both halve and its centroid stays at (19, 0).
((xb[:, None, :] - C[None, :, :]) ** 2).sum(axis=-1) gives a (B, k) distance table, and np.argmin(dist, axis=1) returns the first smallest index.
np.add.at(s, a, xb) sums the batch points into their clusters, and np.bincount(a, minlength=k) counts them.
C[N > 0] = S[N > 0] / N[N > 0][:, None] updates only the clusters with a nonzero count.
Accepts: array
Accepts: array
Accepts: array
Accepts: number
