K-Means Clustering: Algorithm, Choosing k, Limitations

K-means clustering is an unsupervised machine learning algorithm that divides a dataset into k groups. Each point belongs to the cluster whose centroid (the mean of its points) is nearest, and the algorithm minimizes the total squared distance between points and their cluster centroids, called inertia or within-cluster sum of squares.

The method is foundational in data science because it is simple to implement and scales well to large datasets. However, its reliance on distance metrics means that feature scaling and initialization are not just optional tweaks—they are critical to getting a useful result. This page explains the mechanics of the algorithm, how to select the number of clusters, and when you should reach for an alternative.

How the k-means algorithm works

The standard algorithm, often called Lloyd’s algorithm, was proposed by Stuart Lloyd at Bell Labs in 1957 and published in 1982. James MacQueen coined the term “k-means” in 1967. The process follows a strict iterative loop:

Choose the number of clusters, k.
Place k initial centroids within the data space.
Assign every point to its nearest centroid, typically using Euclidean distance.
Move each centroid to the mean (average) of all points assigned to it.
Repeat steps 3 and 4 until the assignments stop changing or a maximum iteration limit is reached.

The algorithm is guaranteed to converge because no step increases the total within-cluster sum of squares. However, it converges to a local optimum, not necessarily the global optimum. The final result depends heavily on where the initial centroids are placed. To mitigate this, practitioners often run the algorithm multiple times with different random starts and select the result with the lowest inertia. This makes the clustering result more stable. Because k-means is unsupervised learning, there are no labels to check against, so inertia is the usual way to compare runs.

A small worked example

Consider a one-dimensional dataset with points: 1, 2, 3, 10, 11, 12. Let k = 2.

Initialization:

Start with centroids at 1 and 12.

Assignment Step:

Points {1, 2, 3} are closer to centroid 1.
Points {10, 11, 12} are closer to centroid 12.

Update Step:

New centroid for the first group: mean(1, 2, 3) = 2.
New centroid for the second group: mean(10, 11, 12) = 11.

Reassignment:

Re-evaluating distances with centroids at 2 and 11 yields the same groups. The algorithm stops.

Inertia Calculation:

The inertia is the sum of squared distances from each point to its centroid: (1 + 0 + 1) + (1 + 0 + 1) = 4.

This example demonstrates why initialization matters. If the initial centroids had been placed poorly, the algorithm might have converged to a suboptimal split.

Initialization and k-means++

Random initialization can lead to poor results if two initial centroids end up close together, leaving other areas of the data space underrepresented. The k-means++ algorithm, introduced by Arthur and Vassilvitskii in 2007, solves this by spreading out the initial centroids.

The process works as follows:

Pick the first centroid randomly from the data points.
For each subsequent centroid, select a point with a probability proportional to its squared distance from the nearest already-chosen centroid.

This ensures that starting points are far apart, leading to faster convergence and better final clusters. This method is now the default initialization in libraries like scikit-learn.

How to choose k

Selecting the number of clusters is one of the hardest parts of using this algorithm. There is no single correct answer, but several methods help guide the decision.

The Elbow Method

Plot the inertia for a range of k values. As k increases, inertia decreases. The “elbow” is the point where the rate of decrease shifts sharply, indicating that adding more clusters yields diminishing returns.

The Silhouette Score

This metric ranges from -1 to 1. A higher score means points are closer to their own cluster than to neighboring clusters. It provides a more nuanced view than the elbow method, especially when clusters are not perfectly spherical.

Practical Constraints

Sometimes the best k is determined by business logic. For example, a marketing team might only be able to act on five distinct customer segments, regardless of what the data suggests.

K-means in Python

Here is a standard implementation using scikit-learn. Note the importance of scaling the data before applying the algorithm.

from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler

X_scaled = StandardScaler().fit_transform(X)

km = KMeans(n_clusters=4, n_init=10, random_state=42)
labels = km.fit_predict(X_scaled)
print(km.inertia_)

Scaling is crucial because k-means uses distance. A feature with large values will dominate the distance calculation, skewing the results. With text data, a common approach is to turn documents into embeddings first and, if needed, reduce their dimensions with principal component analysis. If Euclidean distance does not fit your data, other similarity and correlation measures usually call for a different clustering algorithm.

Limitations and alternatives

K-means is powerful but not universal. Understanding its limitations helps you avoid common pitfalls.

Shape Assumptions: The algorithm assumes roughly spherical clusters of similar size. It struggles with elongated, nested, or irregular shapes. For such data, DBSCAN or Gaussian mixture models are often better choices.
Outlier Sensitivity: Because centroids are means, extreme values can pull them away from the true center of the cluster. K-medoids is less sensitive to outliers.
Fixed k: You must define the number of clusters in advance. If the natural groupings in your data are unknown, this can be a significant constraint.
Data Type: K-means expects numeric features. Purely categorical data requires a different method, such as k-modes.
High Dimensions: In high-dimensional spaces, distance metrics become less meaningful. Reducing dimensions during preprocessing, for example with principal component analysis, often helps.

Where k-means is used

Despite its limitations, this algorithm is widely used in various domains:

Customer Segmentation: Grouping users based on purchasing behavior or demographics.
Image Processing: Color quantization and image compression by reducing the palette of colors.
Document Grouping: Clustering documents based on their embeddings to find related content.
Anomaly Detection: Identifying points that are far from their nearest centroid.
Vector Quantization: Compressing data by mapping high-dimensional vectors to a smaller set of codebook vectors.

FAQ

What is the difference between k-means and k-medoids?

K-means uses the mean of points to define the cluster center, making it sensitive to outliers. K-medoids uses actual data points as centers, which makes it less sensitive to outliers.

Why do I need to scale my data for k-means?

The algorithm relies on distance calculations. If one feature has a much larger range than others, it will dominate the distance metric, causing the algorithm to ignore smaller but potentially important features.

How do I know if my clusters are good?

Use the silhouette score to measure how similar a point is to its own cluster compared to others. You can also visually inspect the clusters if your data has few dimensions, or check if the segments make sense in a business context.

Can k-means handle categorical data?

Not directly. Standard k-means uses Euclidean distance, which doesn’t work well for categorical variables. Use k-modes or convert categories to numeric values using one-hot encoding before applying k-means, though this can increase dimensionality.