Contents

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:

  1. Randomly initialize $k$ points as the cluster centroids.
  2. Calculate the distance from each data point to all centroids.
  3. Assign each point to the nearest cluster centroid.
  4. Compute the mean of the data points assigned to each cluster to find new centroids.
  5. Compare the new centroids with the previous ones. If the difference is small enough, stop the algorithm; otherwise, return to Step 2.

Python Example

https://raw.githubusercontent.com/Josh-test-lab/kmeans-example/refs/heads/main/kmeans_iris_iter/kmeans_iterations.gif
K-means 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.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
# import modules
import numpy as np
import matplotlib.pyplot as plt
import os
from sklearn.datasets import load_iris
from sklearn.metrics import pairwise_distances_argmin

# Iris data
iris = load_iris()  # load data
X = iris.data[:, 2:]  # data
y_true = iris.target  # target
class_names = iris.target_names  # Iris class names

print(f'\nclass names: \n{class_names}')
print(f'\nfeature names: \n{iris.feature_names}')
Execution result reference
1
2
3
4
5
class names: 
['setosa' 'versicolor' 'virginica']

feature names: 
['sepal length (cm)', 'sepal width (cm)', 'petal length (cm)', 'petal width (cm)']

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.

1
2
3
4
5
6
# configs
true_clusters = iris.target_names.shape[0]  # true clusters
n_clusters = 3  # number of clusters we want to classify

output_dir = 'kmeans_iris_iter'  # output directory
os.makedirs(output_dir, exist_ok=True)

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.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
# true centers
true_centers = np.array([X[y_true == i].mean(axis=0) for i in range(true_clusters)])

print(true_centers)

# colors
cmap = plt.get_cmap('tab10')
colors = [cmap(i) for i in range(max(true_clusters, n_clusters))]

# markers
markers = ['o', 's', 'D']
Execution result reference
1
2
3
4
5
6
# colors
cmap = plt.get_cmap('tab10')
colors = [cmap(i) for i in range(max(true_clusters, n_clusters))]

# markers
markers = ['o', 's', 'D']

Ground Truth Scatter Plot

After the preparation steps, we now plot the actual classification of the iris flowers as follows:

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
# plot of true values
plt.figure(figsize=(8, 5))
for i in range(true_clusters):
    # true values
    plt.scatter(X[y_true == i, 0], X[y_true == i, 1], color=colors[i], marker=markers[i], s=20, label=f'{class_names[i]}')
    # true centers
    plt.scatter(true_centers[i, 0], true_centers[i, 1], color=colors[i], marker='X', edgecolor='black', s=120, label=f'center of {class_names[i]}')

plt.title(f'True Clusters')
plt.xlabel('Petal Length (cm)')
plt.ylabel('Petal Width (cm)')
plt.legend(loc='upper left', bbox_to_anchor=(1.05, 1.0), borderaxespad=0)
plt.tight_layout()
plt.savefig(f'{output_dir}/true.png', bbox_inches='tight')
plt.show()
plt.close()

https://raw.githubusercontent.com/Josh-test-lab/kmeans-example/refs/heads/main/kmeans_iris_iter/true.png
The ground truth classification of iris flowers based on petal features.

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.

1
2
3
4
# initialize k-means center points
np.random.seed(123)  # set random seed
init_idx = np.random.choice(len(X), size=n_clusters, replace=False)
centers = X[init_idx]

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.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
# k-means
iteration = 0
while True:
    # assign each point to the nearest center
    labels = pairwise_distances_argmin(X, centers)

    plt.figure(figsize=(8, 5))
    # cluster values
    for i in range(n_clusters):
        plt.scatter(X[labels == i, 0], X[labels == i, 1], facecolors='none', edgecolors=colors[i], linewidths=1.2, s=100, label=f'cluster {i}')

    # true values
    for i in range(true_clusters):
        plt.scatter(X[y_true == i, 0], X[y_true == i, 1], color=colors[i], marker=markers[i], s=20, label=f'{class_names[i]}')
    
    # cluster centers
    for i in range(n_clusters):
        plt.scatter(centers[i, 0], centers[i, 1], color=colors[i], marker='P', edgecolor='black', s=180, label=f'cluster center {i}')

    # true centers
    for i in range(true_clusters):
        plt.scatter(true_centers[i, 0], true_centers[i, 1], color=colors[i], marker='X', edgecolor='black', s=120, label=f'center of {class_names[i]}')

    plt.title(f'K-means Iteration {iteration}')
    plt.xlabel('Petal Length (cm)')
    plt.ylabel('Petal Width (cm)')
    plt.legend(loc='upper left', bbox_to_anchor=(1.05, 1.0), borderaxespad=0)
    plt.tight_layout()
    plt.savefig(f'{output_dir}/kmeans_iter_{iteration:02d}.png', bbox_inches='tight')
    plt.show()
    plt.close()

    # update cluster centers
    new_centers = np.array([X[labels == i].mean(axis=0) if np.any(labels == i) else centers[i] for i in range(n_clusters)])

    # stop condition
    if np.sum((new_centers - centers)**2) < 1e-20:
        break
    
    # next iteration
    centers = new_centers
    iteration += 1

Finally, the entire iterative process is shown as follows:

gallery_made_with_nanogallery2-kmeans

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

References