K-means clustering is a method of vector quantization, originally from signal processing, that partitions n observations into k clusters, where each observation belongs to the cluster with the nearest mean (the cluster center or centroid). This results in a partitioning of the data space into Voronoi cells. The algorithm is widely used in machine learning for unsupervised data analysis, such as customer segmentation, image compression, and pattern recognition.
K-means minimizes within-cluster variances, measured by 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. For instance, better Euclidean solutions can be found using k-medians and k-medoids.
The problem is computationally difficult (NP-hard); however, efficient heuristic algorithms converge quickly to a local optimum. These are usually similar to the expectation–maximization algorithm for mixtures of Gaussian distributions via an iterative refinement approach employed by both k-means and Gaussian mixture modeling. Both use cluster centers to model the data; however, k-means clustering tends to find clusters of comparable spatial extent, while the Gaussian mixture model allows clusters to have different shapes.
The unsupervised k-means algorithm has a loose relationship to the k-nearest neighbor classifier, a popular supervised machine learning technique for classification that is often confused with k-means due to the name. Applying the 1-nearest neighbor classifier to the cluster centers obtained by k-means classifies new data into the existing clusters, known as the nearest centroid classifier or Rocchio algorithm.
Formal Definition
Given a set of observations (x1, x2, ..., xn), where each observation is a d-dimensional real vector, k-means clustering aims to partition the n observations into k (≤ n) sets S = {S1, S2, ..., Sk} so as to minimize the within-cluster sum of squares (WCSS), i.e., variance. Formally, the objective is to find:
argmin over S of the sum over i=1 to k of the sum over x in Si of ||x - μi||^2,
where μi is the mean (also called centroid) of points in Si, and ||·|| is the usual L2 norm. This is equivalent to minimizing the pairwise squared deviations of points in the same cluster, as shown by the identity that the sum of squared distances to the mean equals the average pairwise squared distance.
Algorithm
The most common algorithm, often called Lloyd's algorithm, uses an iterative refinement approach. It starts with an initial set of k centroids, then alternates between two steps: assignment and update. In the assignment step, each observation is assigned to the cluster whose centroid is nearest, typically using Euclidean distance. In the update step, the centroid of each cluster is recalculated as the mean of the assigned points. These steps repeat until assignments no longer change, indicating convergence to a local optimum.
Initialization is crucial; the k-means++ method, which spreads initial centroids, is a popular heuristic to improve the quality of the final clustering. The algorithm is sensitive to the choice of k, and methods like the elbow method or silhouette analysis are used to estimate an appropriate number of clusters.
Properties and Limitations
K-means assumes clusters are spherical and of similar size, which limits its applicability to data with complex cluster shapes. It is also sensitive to outliers, as the mean is influenced by extreme values. The algorithm converges to a local optimum, not necessarily the global one, and different initializations can yield different results. Despite these limitations, its simplicity and scalability make it a popular choice for large datasets, especially in data augmentation and preprocessing pipelines.
Applications
K-means is used in various domains. In artificial intelligence, it serves as a baseline for clustering tasks. In computer vision, it is used for image segmentation and color quantization. In marketing, it helps segment customers based on purchasing behavior. In natural language processing, it can cluster documents or word embeddings. The algorithm is also a building block for more advanced techniques, such as deep learning feature learning and Generative AI models.
Relationship to Other Methods
K-means is related to Gaussian mixture models (GMM), as both use iterative refinement and cluster centers. However, GMM allows clusters to have different shapes and covariances, while k-means assumes isotropic clusters. The nearest centroid classifier, derived from k-means, is a simple supervised classification method. K-means is often confused with k-nearest neighbors (k-NN), but they are distinct: k-means is unsupervised, while k-NN is supervised.
History and Development
The concept of k-means was first proposed by Hugo Steinhaus in 1956, and the term "k-means" was coined by James MacQueen in 1967. Lloyd's algorithm, published in 1957 but not widely known until 1982, is the standard implementation. Over the years, numerous variants have been developed, such as mini-batch k-means for large-scale data and fuzzy c-means for soft clustering.