K-Means Clustering

K-means clustering is an unsupervised learning algorithm that partitions data into k groups, each represented by its centroid, the mean of the points assigned to it. It looks for the partition with the smallest within-cluster sum of squared distances, a quantity called inertia, so points in a cluster sit close to their centroid. The standard procedure is Lloyd's algorithm, which Stuart Lloyd developed at Bell Labs in 1957 for quantizing signals, and the name "k-means" comes from James MacQueen's 1967 paper. Because it is simple, fast, and scales to millions of points, k-means is a default tool for data mining, compression, and exploring the structure of unlabeled data.

How the Algorithm Works#

After choosing starting centroids, Lloyd's algorithm alternates between an assignment step and an update step until the clusters stop changing:

1. Initialization: Choose k starting centroids. Scikit-learn's default, k-means++, picks the first centroid at random and each later one with probability proportional to its squared distance from the nearest centroid already chosen, which spreads the seeds apart.
2. Assignment: Assign every point to its nearest centroid, usually by Euclidean distance.
3. Update: Move each centroid to the mean of the points now assigned to it.
4. Repeat: Return to the assignment step until no point changes cluster or the centroids move less than a tolerance.

Neither step can increase inertia, so the loop always converges, but only to a local minimum that depends on the starting centroids. Finding the global optimum is NP-hard, so implementations run several initializations and keep the one with the lowest inertia. In scikit-learn's KMeans, the default n_init="auto" runs once when seeding with k-means++ and 10 times when seeding randomly.

Choosing the Number of Clusters#

K-means needs k up front. Inertia always falls as k grows, because more centroids sit closer to the points, so the elbow method plots inertia against k and looks for the point where the drop flattens. The silhouette score compares, for each point, the mean distance to its own cluster with the mean distance to the nearest other cluster, and ranges from -1 to 1, with higher values for compact, well-separated clusters. This example scores k from 2 to 6 on data drawn around four centers:

from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs
from sklearn.metrics import silhouette_score

# 600 unlabeled points drawn around 4 hidden centers
X, _ = make_blobs(n_samples=600, centers=4, cluster_std=1.0, random_state=7)

# Fit k-means for several k and score each partition
for k in range(2, 7):
    kmeans = KMeans(n_clusters=k, n_init="auto", random_state=0)
    labels = kmeans.fit_predict(X)
    print(f"k={k}  inertia={kmeans.inertia_:8.0f}  silhouette={silhouette_score(X, labels):.3f}")

Inertia falls from 18,363 at k=2 to 5,241 at k=3 and 1,127 at k=4, then only to 1,039 at k=5 and 922 at k=6, so the elbow sits at k=4. The silhouette score agrees: it peaks at 0.811 for k=4 and drops to 0.687 at k=5, recovering the four centers the data was drawn from.

Assumptions and Variants#

K-means works best when clusters are roughly spherical, similar in size, and separated in Euclidean space. Features on larger scales dominate the distance, so data is usually standardized first. Every point joins some cluster, and a few extreme outliers can drag a centroid away from its group. Several variants relax these limits:

• Mini-batch k-means: Updates centroids from small random batches, trading a little accuracy for much faster training on large datasets.
• K-medoids: Uses actual data points as cluster centers, which makes it more robust to outliers and works with any distance measure.
• Fuzzy c-means: Gives each point a degree of membership in every cluster, in the spirit of fuzzy logic, instead of a single hard assignment.
• Gaussian mixture models: Fit a probability distribution per cluster, which allows elliptical clusters of different sizes and returns soft, probabilistic assignments.

Applications in AI and Computer Vision#

• Color quantization: Clustering an image's pixel colors and replacing each pixel with its centroid reduces a photo with thousands of colors to a palette of k colors, which compresses the image and also gives a fast, rough form of image segmentation.
• Codebooks and vector quantization: The centroids form a codebook, and each vector is stored as the index of its nearest centroid. Classic bag-of-visual-words image retrieval built its vocabularies this way from local image descriptors.
• Anchor box design: YOLOv2 ran k-means on the widths and heights of training boxes, with 1 minus IoU as the distance, to choose its anchor boxes.
• Weight palettization: Exporting YOLO26 to Core ML with quantize=8 clusters the values of each large weight tensor with k-means into 256 centroids, so the exported model stores an 8-bit index per weight plus a small lookup table.
• Customer segmentation: Retailers group customers by purchase frequency, basket size, and recency to tailor offers to each segment.

K-Means vs. Related Algorithms#

• K-nearest neighbors (KNN): The shared "k" causes confusion. KNN is a supervised algorithm that labels a new point from its k closest labeled examples, while k-means is unsupervised and k is the number of clusters it creates.
• DBSCAN: Groups points by density, finds the number of clusters itself, follows clusters of any shape, and labels sparse points as noise instead of forcing them into a cluster.
• Hierarchical clustering: Builds a tree of nested clusters by repeatedly merging the closest groups, so k can be chosen afterward by cutting the tree at any level, at a cost that grows at least quadratically with the number of points.