k-Means Clustering: A Heuristic for an NP-Hard Problem

The objective is easy to state. Partition $n$ observations into $k$ clusters so that each observation belongs to the cluster with the nearest mean, which scikit-learn writes as minimizing the inertia or within-cluster sum-of-squares

${\sum }_{i=0}^n{\min }_{{\mu }_j\in C}\left(\Vert x_i-{\mu }_j{\Vert }^2\right)$

Wikipedia notes what falls out of this geometrically: the result is a partitioning of the data space into Voronoi cells. And it flags a subtlety hiding in the squared term. k-means minimizes within-cluster variances, squared Euclidean distances, “but not regular Euclidean distances, which would be the more difficult Weber problem: the mean optimizes squared errors, whereas only the geometric median minimizes Euclidean distances.”

The idea

The algorithm everyone calls k-means is not a solution to the k-means problem. Wikipedia is unambiguous that the problem is NP-hard, in general Euclidean space even for two clusters, and for a general number of clusters even in the plane. Lloyd’s algorithm is a heuristic that converges quickly to a local optimum, and which local optimum it finds depends entirely on the initialization. Every practical concern with k-means, restarts, k-means++, run-to-run instability, follows from that single gap between the problem and the algorithm.

Lloyd’s algorithm: two steps, alternated

scikit-learn describes the standard procedure in three parts: choose initial centroids, then loop between assigning each sample to its nearest centroid and creating new centroids as the mean of all samples assigned to each previous centroid. Iteration stops when the centroids stop moving significantly, and scikit-learn is precise about its own criterion: “iteration stops when centroids move less than the tolerance,” rather than when the objective’s relative decrease falls below a tolerance.

The assignment step is a Voronoi partition, which Wikipedia states directly: assigning each observation to the nearest mean means “partitioning the observations according to the Voronoi diagram generated by the means.” The update step then moves each centroid to the mean of its region. Both steps can only decrease the objective, which is why the loop terminates.

Two naming facts are worth having straight. Wikipedia records that the term “k-means” was first used by James MacQueen in 1967, though the idea goes back to Hugo Steinhaus in 1956, and that the standard algorithm was first proposed by Stuart Lloyd of Bell Labs in 1957 as a technique for pulse-code modulation, not published as a journal article until 1982. The method was a signal-processing tool before it was a data-science one, developed for vector quantization, and the compression heritage explains the objective’s shape.

Convergence, and the cost of it

Wikipedia gives Lloyd’s running time as $O(nkdi)$ for $n$ vectors of $d$ dimensions, $k$ clusters, and $i$ iterations to convergence, and observes that on data with real clustering structure the iteration count is often small, with results improving only slightly after the first dozen iterations. That makes it “often considered to be of ‘linear’ complexity in practice, although it is in the worst case superpolynomial when performed until convergence.”

The gap between practice and worst case is unusually wide here, and the resolution is a smoothed-analysis result Wikipedia cites: Lloyd’s k-means algorithm has polynomial smoothed running time, so the pathological point sets that force exponential behavior are fragile under small random perturbations and do not appear to arise in practice.

Initialization decides the answer

scikit-learn puts the core caveat plainly: “Given enough time, K-means will always converge, however this may be to a local minimum. This is highly dependent on the initialization of the centroids.” Wikipedia agrees that the algorithm does not guarantee convergence to the global optimum and that the result may depend on the initial clusters, adding the standard mitigation, since the algorithm is usually fast it is common to run it multiple times with different starting conditions.

Naive initializations behave differently in ways that matter. Wikipedia describes Forgy, which randomly chooses $k$ observations from the dataset as the initial means, and Random Partition, which randomly assigns a cluster to every observation and computes initial means from those assignments. “The Forgy method tends to spread the initial means out, while Random Partition places all of them close to the center of the data set.” A comprehensive study by Celebi et al. cited there found that popular methods including Forgy, Random Partition, and Maximin often perform poorly.

k-means++ is the standard fix. scikit-learn describes it as initializing the centroids “to be (generally) distant from each other, leading to probably better results than random initialization.” Wikipedia adds the theoretical property that makes it more than a heuristic-on-a-heuristic: k-means++ “chooses initial centers in a way that gives a provable upper bound on the WCSS objective.” An algorithm with no global guarantee acquires one from its seeding step, which is a satisfying result. scikit-learn also notes k-means++ can be called independently to select seeds for other clustering algorithms.

Inertia is a biased scorer

scikit-learn lists two drawbacks of the objective itself. First, inertia “makes the assumption that clusters are convex and isotropic,” so it responds poorly to elongated clusters or manifolds with irregular shapes. k-means will happily slice a crescent in half. Second, inertia is not normalized, and “in very high-dimensional spaces, Euclidean distances tend to become inflated,” an instance of the curse of dimensionality. The recommended remedy is to run PCA before clustering, which “can alleviate this problem and speed up the computations.” The same curse that breaks kNN breaks k-means, for the same reason: both are built on Euclidean distance.

Choosing k

The algorithm cannot choose $k$; scikit-learn notes flatly that it “requires the number of clusters to be specified.” The two standard approaches both have caveats worth knowing before you reach for them.

The elbow method plots the explained variation as a function of the number of clusters and picks the elbow of the curve. Wikipedia does not stop at describing it. It states that “the notion of an ‘elbow’ is not well-defined and this is known to be unreliable,” citing a 2023 paper by Erich Schubert in ACM SIGKDD Explorations Newsletter titled “Stop using the elbow criterion for k-means and how to choose the number of clusters instead.” The most-taught technique for choosing $k$ has a published argument against it.

Silhouette analysis is the alternative Wikipedia describes: it measures the quality of clustering and gives insight into the separation distance between resulting clusters, where “a higher silhouette score indicates that the object is well matched to its own cluster and poorly matched to neighboring clusters.” That has an actual definition behind it, from Rousseeuw’s 1987 paper, and it scores each point rather than eyeballing a curve.

k-means is EM with the assumptions stripped out

scikit-learn states the connection: “K-means is equivalent to the expectation-maximization algorithm with a small, all-equal, diagonal covariance matrix.” Wikipedia describes the same relationship from the other side, noting that k-means tends to find clusters of comparable spatial extent while a Gaussian mixture model allows clusters to have different shapes. Hard assignment is the limiting case of soft assignment. If your clusters genuinely differ in size, shape, or orientation, the restriction is exactly what is hurting you, and a mixture model is the principled loosening.

Related Notes

• Unsupervised Learning - the no-labels setting clustering belongs to
• PCA and Dimensionality Reduction - the recommended preprocessing before clustering high-dimensional data
• k-Nearest Neighbors - the other Euclidean-distance method with the same dimensionality failure
• Evaluation Metrics - scoring models when there is no ground-truth label to check against
• Features and Representations - why the feature space determines what “nearest” means
• Variance and Covariance - the within-cluster variance the objective minimizes

Sources

• “k-means clustering,” Wikipedia (raw wikitext). https://en.wikipedia.org/w/index.php?title=K-means_clustering&action=raw . Supports partitioning into clusters by nearest mean and the resulting Voronoi cells; minimizing squared Euclidean distances rather than Euclidean distances, with the mean optimizing squared errors and the geometric median minimizing Euclidean distances; NP-hardness in general Euclidean space even for two clusters and for general k even in the plane, with heuristics converging to a local optimum; the assignment step as a Voronoi partition; the naming history (MacQueen 1967, Steinhaus 1956, Lloyd at Bell Labs 1957 for pulse-code modulation, published 1982); the $O(nkdi)$ running time, small iteration counts in practice, and superpolynomial worst case with polynomial smoothed running time; no guarantee of a global optimum and the multiple-restart practice; the Forgy and Random Partition schemes and their spread-out versus centered behavior, plus the Celebi et al. finding; k-means++ giving a provable upper bound on the WCSS objective; the elbow method and the cited 2023 Schubert paper arguing it is unreliable; and silhouette analysis with higher scores meaning better matched to own cluster.
• “Clustering,” scikit-learn User Guide. https://scikit-learn.org/stable/modules/clustering.html . Supports the inertia / within-cluster sum-of-squares objective and its formula; centroids not generally being points of X; the three-step Lloyd formulation and the centroid-movement stopping criterion; inertia assuming convex isotropic clusters and responding poorly to elongated shapes; inertia being unnormalized with Euclidean distances inflated in high dimensions, and PCA before k-means alleviating this; convergence possibly to a local minimum highly dependent on initialization; k-means++ initializing centroids generally distant from each other and being callable independently to seed other algorithms; the requirement that k be specified; and k-means being equivalent to expectation-maximization with a small, all-equal, diagonal covariance matrix.