How does K-Means Clustering work?

Introduction

Clustering is an unsupervised learning task where the goal is to group data points based on similarity, without labeled outcomes. One of the most widely used clustering algorithms is K-means, valued for its simplicity, speed, and interpretability. K-means is often used as an initial exploratory tool to uncover structure in data. Examples include grouping customers by behavior, images by visual similarity, or documents by topic. For an interactive visualization of K-means clustering at work, this website is a good resource to play around with.

What is K-means Clustering?

K-means clustering is a clustering algorithm that partitions a dataset into $K$ clusters, where each data point belongs to the cluster with the nearest centroid (mean). The algorithm runs iteratively:

1. Choose the number of clusters $K$
2. Initialize $K$ centroids
3. Assign each data point to the nearest centroid
4. Recompute centroids as the mean of assigned points
5. Repeat steps 3-4 until convergence

As an example, suppose you have a dataset of 2D points representing customers where:

• The x-axis = annual spending on product A
• The y-axis = annual spending on product B

You suspect that there are three spending patterns (low, medium, high) but the dataset is unlabeled so you have to employ unsupervised learning. Here is how you would use K-means clustering on this dataset in Python:

import numpy as np
from sklearn.cluster import KMeans

# Sample data
X = np.array([
    [1, 2], [1.5, 1.8], [5, 8],
    [8, 8], [1, 0.6], [9, 11]
])

kmeans = KMeans(n_clusters=3, random_state=42)
kmeans.fit(X)

print("Cluster assignments:", kmeans.labels_)
print("Centroids:\n", kmeans.cluster_centers_)

import numpy as np
from sklearn.cluster import KMeans

# Sample data
X = np.array([
    [1, 2], [1.5, 1.8], [5, 8],
    [8, 8], [1, 0.6], [9, 11]
])

kmeans = KMeans(n_clusters=3, random_state=42)
kmeans.fit(X)

print("Cluster assignments:", kmeans.labels_)
print("Centroids:\n", kmeans.cluster_centers_)

This code will produce the following output:

Cluster assignments: [1 1 2 2 1 0]
Centroids:
 [[ 9.         11.        ]
 [ 1.16666667  1.46666667]
 [ 6.5         8.        ]]

This output shows how the six data points were grouped into three clusters by K-means. The cluster assignments array indicates that points labeled $1$ belong to the low-value cluster, points labeled $0$ belong to the mid-range cluster, and points labeled $2$ belong to the high-value cluster. Each centroid represents the mean of the points assigned to that cluster: the low cluster is centered around $(1.17,1.47)$, the mid cluster around $(6.5,8)$, and the high cluster around $(9,11)$. For a visual representation:

What happens here is:

• The three centroids are initialized
• Each point is assigned to the closest centroid, usually via Euclidean distance
• Centroids move the mean of their assigned points
• Repeat the process until the assignments stop changing

Centroid Initialization and Cluster Count

K-means is very sensitive to initialization. Poorly chosen initial centroids can lead to slow convergence, suboptimal local minima, and inconsistent results across runs. To mitigate this, modern implementations use K-means++, which initializes centroids far apart from each other to encourage better cluster separation. Without good initialization, two centroids might start close together, causing one cluster to dominate and another to remain under-utilized.

K-means also requires the number of cluster $K$ to be specified in advance, and this assignment is non-trivial. Common strategies include:

• The Elbow Method: Plot loss (within-cluster variance) vs $K$ and look of a point where improvement slows
• Silhouette Score: Measures how well points fit within their cluster compared to others. Higher is better
• Domain Knowledge: Sometimes the most reliable method. For example, already knowing there are exactly 10 product categories

Loss Function

K-means minimizes within-cluster sum of square distances:

Where:

• $C_i$ is cluster $i$
• $\mu_i$ is the centroid of cluster $i$
• $\left\| x – \mu_i \right\|^2$ is squared Euclidean distance

Each centroid represents the mean of its cluster
Squaring the distance penalizes far-away points heavily
The algorithm seeks compact, tightly packed clusters

Because the loss uses Euclidean distance, K-means implicitly assumes that clusters are roughly spherical, have similar variances, and are separable by distance in the feature space. While K-means does not assume any specific underlying probability distribution, it is highly sensitive to outliers, since extreme points can disproportionately influence centroid locations. As a result, the algorithm performs poorly on elongated, irregular, overlapping, or density-based cluster structures.

Tradeoffs

If these limitations do arise when attempting K-means on a given dataset, alternative clustering approaches are often more appropriate. Density-based methods such as DBSCAN (Density-Based Spatial Clustering of Applications with Noise) can identify arbitrarily shaped clusters and handle outliers naturally, while Gaussian Mixture Models (GMMs) relax the spherical-cluster requirement by modeling clusters with different covariances. For situations where the number of clusters is unknown, techniques such as hierarchical clustering can help determine a suitable cluster structure before committing to a fixed $K$.

Conclusion

K-means clustering is a foundational unsupervised learning algorithm that groups data by minimizing within-cluster variance. It works by iteratively assigning points to centroids and updating those centroids to reflect cluster means. While fast and intuitive, its performance depends heavily on centroid initialization, the choice of $K$, and assumptions about cluster shape. Understanding its loss function helps explain both its strengths and its limitations when applied to complex real-world data.

Video Explanations

• This video, “6.4.7 R6. Segmenting Images – Video 5: K-Means Clustering” by MIT OpenCourseWare provides a tutorial on using the K-means algorithm in R to segment a brain MRI. (Runtime: 7 mins)

• This video, “Machine Learning 13 – K-means” by Stanford Online provides a clear, comprehensive overview of K-means clustering. (Runtime: 20 mins)

Machine Learning 13 – K-means | Stanford CS221: AI (Autumn 2021)

Related Articles

• https://aiml.com/what-is-clustering/
• https://aiml.com/how-does-k-means-work-2/