K-Means Clustering

The cover image was generated by ChatGPT.
Introduction
K-means clustering, often referred to in Chinese as the “k-means algorithm”, is a method used to partition $n$ data points into $k$ clusters. Each point is assigned to the cluster whose center (centroid) is closest to it.
In the context of machine learning, k-means is categorized as an unsupervised learning algorithm. This means it does not require any labeled data during training; instead, it clusters data solely based on the distances between data points. This approach is somewhat similar to how people naturally form small groups in society, without predefined rules, individuals gather based on their similarities.
Principles of K-means
Suppose we have a dataset $$ x_1, x_2, \cdots, x_n $$
Each data point $x_i \in \mathbb{R}^d$, where $i = 1, 2, \cdots, n$. In order to assign the $n$ data points to the nearest of the $k$ clusters, we need to determine, for each point $x_i$, which cluster centroid is closest to it, so that $x_i$ can be assigned to that cluster.
Before that, we must first determine the initial values of the cluster centroids. We randomly select $k$ points as the initial cluster centroids. Then, for each data point $x_i$, we calculate its distance to each centroid and assign it to the cluster with the closest centroid. A variety of distance metrics, such as the Euclidean distance, can be used to calculate distances.
Once every point has been assigned to the nearest cluster centroid, we recalculate the new centroids for each cluster. A new centroid is computed as the mean of all points in the cluster across each dimension. After that, we repeat the steps of reassigning points to clusters and updating the centroids.
Eventually, when the cluster centroids no longer change significantly, or the changes are negligible, the k-means algorithm is considered complete. This process of finding the centroids can be seen as assigning $n$ points into $k$ sets such that the within-cluster sum of squares (WCSS) is minimized. This can be expressed with the following formula:
$$ \argmin_\mathbf{S} \sum_{j=1}^k \sum_{x \in S_j} \| x - \mu_j \|^2, $$
Here, $\mathbf{S}$ denotes the set containing all data points, defined as $\mathbf{S} = \{S_1, S_2, \cdots, S_k\}$, where each $S_j$ represents a cluster for $j = 1, 2, \cdots, k$. The symbol $\mu_j$ denotes the centroid of cluster $S_j$, and $x$ represents all data points within cluster $S_j$. Note that the number of points in each cluster may vary, meaning that $|S_j| \neq |S_m|$.
Algorithm
Simply put, the k-means algorithm can be summarized by the following steps:
- Randomly initialize $k$ points as the cluster centroids.
- Calculate the distance from each data point to all centroids.
- Assign each point to the nearest cluster centroid.
- Compute the mean of the data points assigned to each cluster to find new centroids.
- Compare the new centroids with the previous ones. If the difference is small enough, stop the algorithm; otherwise, return to Step 2.
Python Example
Dataset Description
We will use the Iris dataset as an example. The Iris dataset is a classic dataset frequently used in machine learning and statistical analysis, especially for classification and visualization tasks. The goal is to predict the species of iris flowers.
This dataset includes:
- Number of samples: 150
- Number of features: 4 numerical features
- Sepal length (in centimeters)
- Sepal width (in centimeters)
- Petal length (in centimeters)
- Petal width (in centimeters)
- Number of classes: 3 species of iris flowers
- Setosa
- Versicolor
- Virginica
Setup
In Python, we can load the Iris dataset using the sklearn module. In the following example, we use the last two features, petal length and petal width, as the input for classification and visualization.
| |
| |
Next, we specify the number of clusters. Since we already know that there are three species of iris flowers, we directly set the ground truth cluster count true_clusters based on the dataset’s shape, and define n_clusters as 3 for our clustering task. If you want to experiment with more clusters, you can also modify the value of n_clusters.
| |
We now define the ground truth cluster centers to facilitate comparison with the clustering results of k-means. At the same time, we assign different colors to each class to make the visual comparison easier later on.
| |
| |
Ground Truth Scatter Plot
After the preparation steps, we now plot the actual classification of the iris flowers as follows:
| |

As shown in the figure above, blue represents Setosa, orange represents Versicolor, and green represents Virginica; while ✕ indicates the cluster centroids.
K-means
First, initialize the cluster centroids. Here, we fix the random seed to 123 and randomly select n_clusters points as the initial cluster centroids.
| |
Next, we proceed with iterative training to find the cluster centroids. During each iteration, we will plot the results to visualize the movement of the cluster centers, while also displaying the ground truth on the plots for comparison.
| |
Finally, the entire iterative process is shown as follows:
In the figure above, solid dots represent the actual classifications of the iris species, while hollow circles (◯) indicate the clusters assigned by the algorithm; ✕ marks the true centroids of each species, and ✛ denotes the centroids of the clusters found by k-means.
We can observe that, during each iteration, the cluster centroids found by the algorithm (✛) gradually approach the true centroids (✕). Although a few misclassifications still occur near the boundary areas, overall, k-means correctly assigns most data points to their respective clusters.
Conclusion
The k-means algorithm is a simple yet efficient clustering method widely used in various fields such as image processing, market segmentation, and bioinformatics. By iteratively updating cluster centroids and reassigning data points, k-means can automatically discover underlying structures and patterns without the need for labeled data.
Although k-means performs well in many scenarios, it also has some limitations, such as sensitivity to the initial centroids, being applicable only to convex-shaped clusters, and vulnerability to outliers. Therefore, in practical applications, it is important to carefully select the algorithm based on the characteristics of the data or consider combining it with other methods to achieve more stable and accurate clustering results.
Environment
- Operating System: Windows 11 24H2
- Programming Language: Python 3.12.9
Further Learning
- The ipynb file used in this article.
References
k-平均演算法. (April 21, 2025). Wikipedia, The Free Encyclopedia. Retrieved July 11, 2025, from https://zh.wikipedia.org/zh-tw/K-平均算法
Alex lin. (January 1, 2022). 【機器學習筆記】聚類分析K-means clustering. Medium. Retrieved July 11, 2025, from https://medium.com/@SCU.Datascientist/python學習筆記-聚類分析k-means-clustering-63fd65027c98
Jason Chen. (July 18, 2019). 【機器學習】聚類分析 K-means Clustering. Jason Chen’s Blog. Retrieved July 11, 2025, from https://jason-chen-1992.weebly.com/home/-k-means-clustering
ramonliao. (November 9, 2018). [演算法] K-means 分群 (K-means Clustering). iT 邦幫忙. Retrieved July 11, 2025, from https://ithelp.ithome.com.tw/articles/10209058
Tommy Huang. (April 27, 2018). 機器學習: 集群分析 K-means Clustering. Medium. Retrieved July 11, 2025, from https://chih-sheng-huang821.medium.com/機器學習-集群分析-k-means-clustering-e608a7fe1b43










