K-Means Clustering: Algorithm, Math, and Implementation

Master K-Means Clustering. Learn the mathematics behind centroid optimization, the Elbow method, Silhouette score, and practical Python implementation.

Introduction

Imagine you have a messy pile of clothes of different sizes (S, M, L, XL). You want to group them into four distinct piles, but there are no size tags on them. You would likely start by picking four random clothes as representatives, then group all other clothes by comparing their sizes to these four representatives, and gradually adjust your piles.

This is exactly how K-Means Clustering works. It is one of the simplest, most popular, and powerful unsupervised machine learning algorithms used to partition data points into $K$ distinct, non-overlapping groups (clusters).

What You Will Learn

  • How the K-Means algorithm works step-by-step.
  • The mathematics behind cluster assignment and centroid updates.
  • How to choose the optimal number of clusters ($K$) using the Elbow Method and Silhouette Score.
  • How to implement K-Means from scratch and using scikit-learn in Python.

Why This Topic Matters

K-Means is a foundational clustering algorithm used across industries. Whether you are building a customer segmentation model for a marketing campaign, compressing image sizes, or grouping search engine results, K-Means is usually the first algorithm you deploy because of its speed and ease of interpretation.

Prerequisites

Detailed Explanation

The objective of K-Means is to group $N$ data points into $K$ clusters such that points within the same cluster are as close to each other as possible, and as far from points in other clusters as possible.

Mathematical Formulation

K-Means minimizes the sum of squared distances between data points and their corresponding cluster centroids. This is known as Inertia or the Within-Cluster Sum of Squares (WCSS):

$$WCSS = \sum_{j=1}^{K} \sum_{i \in S_j} ||x_i - \mu_j||^2$$

Where:

  • $K$ is the number of clusters.
  • $S_j$ is the set of data points in the $j$-th cluster.
  • $x_i$ is a data point in that cluster.
  • $\mu_j$ is the centroid (mean) of the $j$-th cluster.
  • $||\cdot||$ denotes the Euclidean distance.

Step-by-Step Algorithm

  1. Initialization: Choose $K$ random data points to act as initial centroids (cluster centers).
  2. Assignment Step: Assign each data point to the nearest centroid based on Euclidean distance: $$d(x, y) = \sqrt{\sum_{m=1}^{d} (x_m - y_m)^2}$$
  3. Update Step: Calculate the new centroids by taking the average (mean) of all data points assigned to each cluster: $$\mu_j = \frac{1}{|S_j|} \sum_{i \in S_j} x_i$$
  4. Convergence: Repeat steps 2 and 3 until the centroids no longer move, or the maximum number of iterations is reached.

How to Choose the Optimal $K$?

Choosing the correct $K$ is crucial. Two main techniques are used:

1. The Elbow Method

Plot the WCSS (Inertia) against different values of $K$. As $K$ increases, WCSS naturally drops because the clusters get smaller. The point where the rate of decrease shifts dramatically (creating an "elbow" shape) represents the optimal $K$.

2. Silhouette Coefficient

Measures how similar a data point is to its own cluster compared to other clusters. The score ranges from $-1$ to $+1$:

  • Near +1: The point is well clustered.
  • Near 0: The point is on the boundary between two clusters.
  • Near -1: The point is in the wrong cluster.

Visual Diagram (Mermaid)

graph TD
    A[Start: Choose K Centroids] --> B[Assign Points to Closest Centroid]
    B --> C[Compute New Centroids Average of Points]
    C --> D{Did Centroids Move?}
    D -->|Yes| B
    D -->|No/Converged| E[Final Clusters Found]
    
    style A fill:#3B82F6,stroke:#fff,color:#fff
    style E fill:#10B981,stroke:#fff,color:#fff

Python Code Examples

Here is how to implement K-Means clustering and find the optimal $K$ using the Elbow method.

import numpy as np
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans
from sklearn.datasets import make_blobs

# 1. Generate synthetic data (3 distinct groups)
X, y = make_blobs(n_samples=300, centers=3, cluster_std=0.60, random_state=42)

# 2. Elbow Method to find optimal K
wcss = []
k_range = range(1, 11)

for k in k_range:
    kmeans = KMeans(n_clusters=k, init='k-means++', max_iter=300, n_init=10, random_state=42)
    kmeans.fit(X)
    wcss.append(kmeans.inertia_)

# Print WCSS values
for k, val in zip(k_range, wcss):
    print(f"K = {k}: WCSS = {round(val, 2)}")

# 3. Fit K-Means with optimal K=3
optimal_kmeans = KMeans(n_clusters=3, init='k-means++', max_iter=300, n_init=10, random_state=42)
y_pred = optimal_kmeans.fit_predict(X)

print("\nCentroid Coordinates:")
print(optimal_kmeans.cluster_centers_)

Industry Use Cases

  • E-commerce Customer Segmentation: Grouping customers based on purchase history, frequency, and spending score to design tailored marketing strategies.
  • Image Compression: Reducing the number of unique colors in an image by clustering pixels by their RGB values and replacing them with centroid values.
  • Document Clustering: Grouping news articles into topics (sports, tech, politics) based on word frequencies.

Advantages & Limitations

Advantages

  • Simple to understand and implement.
  • Computationally efficient with $O(T \cdot K \cdot N \cdot D)$ time complexity, making it scale well to large datasets.
  • Guaranteed convergence to a local optimum.

Limitations

  • Sensitive to Initialization: Poor starting centroids can lead to suboptimal clusters ( mitigated using the k-means++ initialization method).
  • Requires Predefining $K$: You must guess or calculate the number of clusters beforehand.
  • Sensitive to Outliers: Since centroids are means, outliers can skew cluster boundaries.
  • Assumes Spherical Clusters: Fails on complex shapes (e.g., moons, rings).

FAQs

Q: What is the difference between k-means and k-means++? A: Standard K-Means picks initial centroids purely at random. k-means++ chooses initial centroids that are spaced far apart from one another, which leads to faster convergence and better cluster qualities.

Q: Can K-Means handle categorical data? A: No, because you cannot calculate the mathematical mean (centroid) of text categories. For categorical data, you should use algorithms like K-Modes or convert features into numeric representations.

Summary

K-Means is a centroid-based clustering algorithm that groups data points by minimizing the distance between points and their respective cluster centers. While it is fast, scalable, and simple, you must be careful to choose the right number of clusters ($K$) using the Elbow Method or Silhouette Score, and keep in mind its limitations with non-spherical shapes and outliers.

Next Topic

What if you want to cluster data without specifying the number of clusters beforehand, and instead see how they merge hierarchically? Let's explore: Hierarchical Clustering.

← Back to Syllabus