K-Means Update Step

Implement the update step of the K-Means algorithm.

Given data points, current assignments, and the number of clusters $k$ , compute new centroids as the mean of all points assigned to each cluster.

Return the new centroids as a 2D list. If a cluster has no points, keep its centroid at the origin: a list of $d$ integer zeros (e.g. [0, 0] for 2D points, not [0.0, 0.0]).

Round each centroid coordinate to 4 decimal places.

Example

Input

X = [[1, 0], [3, 0], [8, 0], [10, 0]]
assignments = [0, 0, 1, 1]
k = 2

Output

[[2.0, 0.0], [9.0, 0.0]]

Reasoning

X holds $n = 4$ points in $d = 2$ dimensions, all lying on the horizontal axis (every second coordinate is $0$ ). assignments labels each point with its cluster: the first two belong to cluster $0$ , the last two to cluster $1$ . With k $= 2$ the result must contain exactly two centroids.
The new centroid of cluster $c$ is the coordinate-wise mean of its members: $\mu_c = \frac{1}{|S_c|}\sum_{x \in S_c} x$ , where $S_c$ is the set of points labelled $c$ .
Cluster $0$ has $S_0 = \{(1,0), (3,0)\}$ , so $\mu_0 = \left(\frac{1+3}{2}, \frac{0+0}{2}\right) = (2, 0)$ .
Cluster $1$ has $S_1 = \{(8,0), (10,0)\}$ , so $\mu_1 = \left(\frac{8+10}{2}, \frac{0+0}{2}\right) = (9, 0)$ .
Rounded to 4 decimal places these stay [[2.0, 0.0], [9.0, 0.0]]; they print as decimals because a mean is a division.
Notice each centroid sits exactly midway between its two points, and the old centroid positions play no role. Had no point been labelled $1$ , $\mu_1$ would be the origin instead.

Constraints

X: 2D list of data points (n x d)
assignments: list of n cluster indices (0 to k-1)
k: number of clusters
Return 2D list of k centroids, each rounded to 4 decimal places
Empty clusters get centroid at origin

Learn this in

ClusteringML Study Plan

More from Machine Learning 2

Kernel PCAhard PCA via Covariance Matrixmedium Variance Explained Ratioeasy Feature Scalingeasy GAN Discriminator Losshard

K-Means Update Step

Example

Constraints

Learn this in

More from Machine Learning 2

Background Knowledge

Continue the full explanation

Editor locked