K-Means Assignment Step

Implement the assignment step of the K-Means algorithm.

Given data points and current centroids, assign each point to the nearest centroid using Euclidean distance.

Return a list of cluster assignments (0-indexed centroid indices). If a point is equidistant from multiple centroids, assign it to the one with the smallest index.

Example

Input

X = [[1, 0], [2, 0], [8, 0], [9, 0]]
centroids = [[1.5, 0], [8.5, 0]]

Output

[0, 0, 1, 1]

Reasoning

X holds $n = 4$ points in $d = 2$ dimensions, all on the horizontal axis; centroids holds $k = 2$ cluster centres, at $(1.5, 0)$ and $(8.5, 0)$ . The answer gives, for each point, the index of its closest centre.
The Euclidean distance between a point $x$ and a centre $c$ is $\lVert x - c \rVert = \sqrt{\sum_{j=1}^{d} (x_j - c_j)^2}$ .
Point $(1, 0)$ : distance to centre 0 is $\sqrt{(1-1.5)^2 + 0^2} = 0.5$ , distance to centre 1 is $\sqrt{(1-8.5)^2} = 7.5$ . Centre 0 is closer, so its label is $0$ .
Point $(2, 0)$ : distances $0.5$ and $6.5$ , label $0$ . Point $(8, 0)$ : distances $6.5$ and $0.5$ , label $1$ . Point $(9, 0)$ : distances $7.5$ and $0.5$ , label $1$ .
Collecting the labels in the order the points appear gives [0, 0, 1, 1]: a list of $n$ integers, one per point, each between $0$ and $k-1$ .
Notice the tie rule: a point exactly halfway between two centres, such as $5$ between $0$ and $10$ (both at distance $5$ ), takes the smaller index, $0$ .

Constraints

X: 2D list of data points (n x d)
centroids: 2D list of centroid positions (k x d)
Return list of n integers (cluster assignments, 0-indexed)

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

Concept primer

Background Knowledge

The K-Means algorithm is a widely used unsupervised learning technique for clustering data points into K distinct groups, or clusters. The goal of K-Means is to identify patterns or structures in the data by grouping similar points together. The algorithm consists of two main steps: assignment and update. In the assignment step, each data point is assigned to the nearest centroid, which is the center of a cluster. The update step involves recalculating the centroids based on the assigned points.

The assignment step relies on calculating the Euclidean distance between each data point and the current centroids. The Euclidean distance between two points $x = (x_1, x_2, \ldots, x_n)$ and $y = (y_1, y_2, \ldots, y_n)$ is given by $\sqrt{(x_1 - y_1)^2 + (x_2 - y_2)^2 + \ldots + (x_n - y_n)^2}$ . This distance metric is used to determine the nearest centroid for each point. If a point is equidistant from multiple centroids, it is assigned to the one with the smallest index, ensuring a consistent assignment.

The K-Means algorithm is an iterative process, with the assignment and update steps alternating until convergence or a stopping criterion is reached. Understanding the assignment step is crucial for implementing the K-Means algorithm, as it forms the foundation for the overall clustering process.

Algorithm/Approach

The general approach to solving this problem involves iterating over each data point and calculating its Euclidean distance to each centroid. The point is then assigned to the centroid with the minimum distance. If multiple centroids have the same minimum distance, the point is assigned to the one with the smallest index. This process is repeated for all data points, resulting in a list of cluster assignments.

Step-by-Step Strategy

To implement the assignment step, follow these steps:

Iterate over each data point
For each data point, calculate the Euclidean distance to each centroid
Identify the centroid with the minimum distance for each point
If multiple centroids have the same minimum distance, assign the point to the one with the smallest index
Store the cluster assignment for each point

Common Pitfalls

When implementing the assignment step, watch out for:

Incorrectly calculating the Euclidean distance between points and centroids
Failing to handle cases where a point is equidistant from multiple centroids
Not storing the cluster assignments correctly

Time & Space Complexity

The expected time complexity for the assignment step is $O(n \cdot K \cdot d)$ , where $n$ is the number of data points, $K$ is the number of centroids, and $d$ is the dimensionality of the data. The space complexity is $O(n)$ , as we need to store the cluster assignment for each point.

solution.py

Python 3.13

Test results