Skip to content

K-means Algorithm

K-Means is one of the most popular and widely used unsupervised clustering algorithms. Its goal is to partition a dataset into a pre-defined number (K) of distinct, non-overlapping clusters.

It aims to minimize the within-cluster variance, also known as inertia or Within-Cluster Sum of Squares (WCSS). This means it tries to make the data points within each cluster as similar (close) to each other as possible, while keeping the clusters themselves as distinct (far apart) as possible.

Each cluster is represented by its center, called the centroid, which is typically the mean of all data points belonging to that cluster.

How K-Means Works: The Expectation-Maximization Approach

Section titled “How K-Means Works: The Expectation-Maximization Approach”

K-Means uses an iterative refinement approach that alternates between two steps (often viewed as a form of Expectation-Maximization):

  1. Initialization:
  2. a. Specify K: Choose the desired number of clusters, K.
  3. b. Initialize Centroids: Select K initial centroids. Common methods include random selection of K data points or the more sophisticated ‘k-means++’ initialization (default in Scikit-learn), which tries to place initial centroids far apart.
  4. Iteration (Repeat until convergence):
  5. a. Expectation Step (Assignment): Assign each data point to the cluster whose centroid is nearest (usually based on Euclidean distance).
  6. b. Maximization Step (Update): Recalculate the position of the centroid for each cluster by taking the mean of all data points assigned to that cluster in the previous step.
  7. Convergence: The algorithm stops when the assignments no longer change, the centroids stabilize, or a maximum number of iterations is reached.

Unlike hierarchical clustering, K-Means requires specifying K beforehand. Choosing the optimal K is often subjective and depends on the specific application and data. Common methods include:

  • Elbow Method: Plot the inertia (WCSS) for different values of K. Look for an ‘elbow’ point where the rate of decrease in inertia slows down significantly. This suggests adding more clusters beyond that point yields diminishing returns.
  • Silhouette Score: Measures how similar a point is to its own cluster compared to other clusters. Values range from -1 to +1. Higher values indicate better-defined clusters. Calculate the average silhouette score for different K values and look for a peak.
  • Domain Knowledge: Prior knowledge about the data might suggest a natural number of groups.
  • Feature Scaling: Like KNN, K-Means relies on distance calculations (typically Euclidean). Therefore, feature scaling (e.g., Standardization) is crucial to prevent features with larger ranges from dominating the clustering.
  • Initialization Sensitivity: The final clustering result can depend on the initial placement of centroids. Running the algorithm multiple times with different random initializations (n_init parameter in Scikit-learn) and choosing the best result (lowest inertia) is standard practice. ‘k-means++’ initialization helps mitigate this.
  • Assumptions: K-Means implicitly assumes clusters are spherical, roughly equal in size, and have similar density. It may struggle with clusters that are elongated, non-convex, or have vastly different sizes/densities.

Let’s apply K-Means to simple synthetic data with known cluster structures.

import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from sklearn.datasets import make_blobs
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
# 1. Generate synthetic data (e.g., 4 blobs)
X, y_true = make_blobs(n_samples=400, centers=4, cluster_std=0.80, random_state=42)
# 2. Scale the data (important for K-Means)
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
# 3. Apply K-Means
# n_clusters: The desired number of clusters (K)
# init: Initialization method ('k-means++' is default and recommended)
# n_init: Number of times to run with different centroid seeds ('auto' is new default, often 10)
# max_iter: Maximum iterations per run
# random_state: For reproducibility of centroid initialization
kmeans = KMeans(n_clusters=4, init='k-means++', n_init='auto', random_state=42)
# Fit the model and predict cluster labels
# fit_predict combines fitting and getting labels
labels = kmeans.fit_predict(X_scaled)
# Get cluster centroids (in scaled space)
centroids_scaled = kmeans.cluster_centers_
# Get inertia (WCSS)
inertia = kmeans.inertia_
print(f"K-Means Inertia (WCSS): {inertia:.2f}")
# 4. Visualize the results
plt.figure(figsize=(10, 7))
# Plot data points colored by assigned cluster label
plt.scatter(X_scaled[:, 0], X_scaled[:, 1], c=labels, s=50, cmap='viridis', alpha=0.7)
# Plot the centroids
plt.scatter(centroids_scaled[:, 0], centroids_scaled[:, 1], c='red', s=200, marker='X', label='Centroids')
plt.title('K-Means Clustering Results (k=4)')
plt.xlabel('Feature 1 (Scaled)')
plt.ylabel('Feature 2 (Scaled)')
plt.legend()
plt.grid(True)
plt.show()
K-Means Inertia (WCSS): 255.33

The output gives the inertia score (lower is generally better for a given K). The plot shows the data points colored by their assigned cluster and the final positions of the cluster centroids (red ‘X’ marks). K-Means successfully identified the four distinct groups in this synthetic data.

Let’s apply K-Means to the Digits dataset to see if it can group similar handwritten digits without using the true labels.

import matplotlib.pyplot as plt
import numpy as np
from sklearn.datasets import load_digits
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import confusion_matrix, adjusted_rand_score, silhouette_score
import seaborn as sns
# 1. Load Digits dataset
digits = load_digits()
X = digits.data
y_true = digits.target
# The dataset has 1797 samples, each being an 8x8 image flattened to 64 features.
print(f"Digits data shape: {X.shape}")
# 2. Scale the data
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X)
# 3. Apply K-Means (We know there are 10 digits, so set k=10)
kmeans_digits = KMeans(n_clusters=10, init='k-means++', n_init='auto', random_state=42)
labels = kmeans_digits.fit_predict(X_scaled)
# 4. Evaluate Clustering (using metrics that don't require true labels first)
sil_score = silhouette_score(X_scaled, labels)
print(f"\nSilhouette Score: {sil_score:.4f}")
# Evaluate using true labels (for analysis, not for guiding clustering itself)
ari_score = adjusted_rand_score(y_true, labels)
print(f"Adjusted Rand Index (ARI): {ari_score:.4f}")
# Optional: Build a confusion matrix mapping cluster labels to true digit labels
# This requires finding the most frequent true label in each cluster
from scipy.stats import mode
conf_mat_km = np.zeros((10, 10), dtype=int)
for i in range(10): # Iterate through clusters 0-9
mask = (labels == i)
# Find the most common true label within this cluster
cluster_true_labels = y_true[mask]
if len(cluster_true_labels) > 0:
true_label_mode = mode(cluster_true_labels)[0]
# Count how many of each true label fall into this cluster
for j in range(10):
conf_mat_km[true_label_mode, i] = np.sum(cluster_true_labels == j)
plt.figure(figsize=(8, 6))
sns.heatmap(conf_mat_km, annot=True, fmt='d', cmap='Blues')
plt.xlabel('Predicted Cluster Index')
plt.ylabel('Dominant True Label in Cluster')
plt.title('K-Means Clustering vs True Labels (Digits)')
plt.show()
# 5. Visualize Centroids as Images
centroids_scaled = kmeans_digits.cluster_centers_
# Inverse transform centroids back to original scale for visualization
centroids = scaler.inverse_transform(centroids_scaled)
fig, ax = plt.subplots(2, 5, figsize=(8, 3))
fig.suptitle('K-Means Cluster Centroids (as Digit Images)', fontsize=14)
centers = centroids.reshape(10, 8, 8)
for axi, center in zip(ax.flat, centers):
axi.set(xticks=[], yticks=[])
axi.imshow(center, interpolation='nearest', cmap=plt.cm.binary)
plt.show()
Digits data shape: (1797, 64)
Silhouette Score: 0.1862
Adjusted Rand Index (ARI): 0.6687

The Silhouette Score and ARI indicate reasonably good clustering quality, although not perfect. The confusion matrix (if implemented fully) would show how well the clusters align with the true digit labels. The visualization of the centroids shows that K-Means has learned recognizable (though sometimes blurry) representations of the digits 0-9 based purely on the pixel data.

  • Simplicity and Speed: Relatively easy to understand and implement. Computationally faster than hierarchical clustering, especially for large datasets (typically O(NKI*D) where N=samples, K=clusters, I=iterations, D=dimensions).
  • Scalability: Works well for large datasets where other methods might be too computationally expensive.
  • Convergence: Guaranteed to converge (though possibly to a local optimum).
  • Need to Specify K: Requires the number of clusters (K) to be specified beforehand.
  • Sensitivity to Initialization: Results can vary depending on the initial placement of centroids (mitigated by multiple runs n_init and ‘k-means++’).
  • Sensitivity to Outliers: Centroids (based on means) are sensitive to outliers, which can pull them away from the actual cluster centers.
  • Assumes Spherical Clusters: Struggles with clusters that are non-spherical, elongated, or have complex shapes.
  • Assumes Equal Variance/Size: Implicitly assumes clusters have similar variance and number of points. May perform poorly on clusters of vastly different sizes or densities.

K-Means is widely used for:

  • Customer Segmentation: Grouping customers based on purchasing behavior, demographics, etc.
  • Document Clustering: Grouping similar documents based on topics or content.
  • Image Segmentation: Grouping pixels with similar colors or intensities to identify objects or regions.
  • Image Compression (Vector Quantization): Reducing the number of colors in an image by clustering colors and replacing pixels with their cluster centroid color.
  • Anomaly Detection: Points far from any cluster centroid might be considered anomalies.
  • Data Preprocessing: Can be used to create features for supervised learning (e.g., cluster ID as a feature).