K-Means Clustering Algorithm

K-Means partitions n observations into k clusters, where each observation belongs to the cluster with the nearest mean (centroid). It is the most widely used clustering algorithm — simple, fast, and scalable.

The Algorithm

INPUT:  dataset D, number of clusters k

1. INITIALISE:  choose k initial centroids (randomly, or by k-means++)

2. REPEAT until convergence:

   a) ASSIGNMENT STEP:
      assign each data point to the cluster whose centroid is NEAREST
      (usually by Euclidean distance)

   b) UPDATE STEP:
      recompute each centroid as the MEAN of all points now assigned to it

3. STOP when centroids no longer move (or move less than a tolerance,
   or a maximum iteration count is reached)

OUTPUT: k clusters with their centroids

Objective Function

K-Means minimises the Within-Cluster Sum of Squares (WCSS), also called inertia:

              k
   WCSS  =   Σ    Σ    || x − μᵢ ||²
             i=1  x∈Cᵢ

   Cᵢ = the i-th cluster,  μᵢ = its centroid

Each iteration is guaranteed to reduce (or keep constant) the WCSS, which is why the algorithm always converges — though possibly to a local minimum, not the global one.

Worked Example — By Hand

Data points: A(2,10), B(2,5), C(8,4), D(5,8), E(7,5), F(6,4), G(1,2), H(4,9) Let k = 3 with initial centroids μ₁ = A(2,10), μ₂ = D(5,8), μ₃ = G(1,2).

Iteration 1 — Assignment step (Euclidean distances):

Sample calculation for B(2,5) to μ₂(5,8):
   d = √[(2−5)² + (5−8)²] = √[9 + 9] = √18 = 4.243

Cluster 1: {A}
Cluster 2: {C, D, E, F, H}
Cluster 3: {B, G}

Iteration 1 — Update step:

μ₁ = mean{A}          = (2, 10)

μ₂ = mean{C,D,E,F,H}  = ((8+5+7+6+4)/5, (4+8+5+4+9)/5)
                      = (30/5, 30/5) = (6, 6)

μ₃ = mean{B,G}        = ((2+1)/2, (5+2)/2) = (1.5, 3.5)

Iteration 2 — Assignment step with new centroids μ₁(2,10), μ₂(6,6), μ₃(1.5,3.5):

Cluster 1: {A, H}          <- H moved from cluster 2 to cluster 1
Cluster 2: {C, D, E, F}
Cluster 3: {B, G}

New centroids:
   μ₁ = ((2+4)/2, (10+9)/2) = (3, 9.5)
   μ₂ = ((8+5+7+6)/4, (4+8+5+4)/4) = (6.5, 5.25)
   μ₃ = (1.5, 3.5)              [unchanged]

Iteration 3 produces no further reassignments → converged.

FINAL RESULT:
   Cluster 1: {A(2,10), H(4,9)}          centroid (3, 9.5)
   Cluster 2: {C(8,4), D(5,8), E(7,5), F(6,4)}   centroid (6.5, 5.25)
   Cluster 3: {B(2,5), G(1,2)}           centroid (1.5, 3.5)

Choosing k — The Elbow Method

Plot WCSS against k. WCSS always falls as k rises (at k = n, WCSS = 0), so look for the elbow — the point where the rate of decrease sharply flattens. That is the k beyond which extra clusters add little.

   WCSS
     │●
     │ ●
     │  ●
     │   ●___  <- ELBOW at k=3: adding more clusters barely helps
     │       ●___●___●___●
     └────────────────────── k
      1  2  3  4  5  6  7

Alternatives: the silhouette method (pick k maximising the average silhouette score), the gap statistic, and plain domain knowledge ("marketing wants exactly 4 segments").

Python

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.cluster import KMeans

# Reproducing the hand-worked example
points = np.array([[2,10],[2,5],[8,4],[5,8],[7,5],[6,4],[1,2],[4,9]])
labels_txt = list("ABCDEFGH")

km = KMeans(n_clusters=3, init="random", n_init=1,
            random_state=1, max_iter=100)
clusters = km.fit_predict(points)

for lbl, pt, c in zip(labels_txt, points, clusters):
    print(f"{lbl}{tuple(pt)} -> cluster {c}")
print("\nCentroids:\n", km.cluster_centers_.round(2))
print("Inertia (WCSS):", round(km.inertia_, 3))
print("Iterations to converge:", km.n_iter_)

# THE ELBOW METHOD on a realistic customer dataset
from sklearn.datasets import make_blobs
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import silhouette_score

X, _ = make_blobs(n_samples=400, centers=4, cluster_std=1.1, random_state=42)
X = StandardScaler().fit_transform(X)

wcss, silhouettes = [], []
k_range = range(2, 11)
for k in k_range:
    km = KMeans(n_clusters=k, n_init=10, random_state=42).fit(X)
    wcss.append(km.inertia_)
    silhouettes.append(silhouette_score(X, km.labels_))

fig, axes = plt.subplots(1, 2, figsize=(14, 4.5))

axes[0].plot(k_range, wcss, marker="o", color="#168B99", linewidth=2)
axes[0].set_xlabel("k"); axes[0].set_ylabel("WCSS (inertia)")
axes[0].set_title("Elbow Method"); axes[0].grid(alpha=0.3)

axes[1].plot(k_range, silhouettes, marker="s", color="#10b981", linewidth=2)
best_k = list(k_range)[int(np.argmax(silhouettes))]
axes[1].axvline(best_k, color="#ef4444", linestyle="--", label=f"best k = {best_k}")
axes[1].set_xlabel("k"); axes[1].set_ylabel("Silhouette score")
axes[1].set_title("Silhouette Method"); axes[1].legend(); axes[1].grid(alpha=0.3)

plt.tight_layout(); plt.show()
print(f"Elbow and silhouette both point to k = {best_k}")

# CUSTOMER SEGMENTATION — the flagship business application
np.random.seed(42)
customers = pd.DataFrame({
    "age":            np.concatenate([np.random.normal(25, 4, 60),
                                      np.random.normal(45, 6, 60),
                                      np.random.normal(38, 5, 60)]),
    "annual_income":  np.concatenate([np.random.normal(30000, 6000, 60),
                                      np.random.normal(90000, 15000, 60),
                                      np.random.normal(60000, 9000, 60)]),
    "spending_score": np.concatenate([np.random.normal(75, 10, 60),
                                      np.random.normal(35, 10, 60),
                                      np.random.normal(55, 10, 60)]),
})

features = ["age", "annual_income", "spending_score"]
X_scaled = StandardScaler().fit_transform(customers[features])   # scaling is essential

km = KMeans(n_clusters=3, n_init=10, random_state=42)
customers["segment"] = km.fit_predict(X_scaled)

profile = customers.groupby("segment")[features].mean().round(1)
profile["size"] = customers["segment"].value_counts().sort_index()
print(profile)
#            age  annual_income  spending_score  size
# segment
# 0         25.1        30169.5            74.6    60
# 1         44.9        89632.2            34.6    60
# 2         38.2        60113.9            55.3    60

print(f"\nSilhouette score: {silhouette_score(X_scaled, customers['segment']):.4f}")

# Business naming of the discovered segments
names = {0: "Young Big Spenders — target with trendy, affordable products",
         1: "Affluent but Frugal — target with premium value propositions",
         2: "Balanced Mid-Market — target with loyalty programmes"}
for seg, desc in names.items():
    print(f"Segment {seg}: {desc}")

# Visualising the segments
fig, axes = plt.subplots(1, 2, figsize=(14, 5))

sc = axes[0].scatter(customers["annual_income"], customers["spending_score"],
                     c=customers["segment"], cmap="viridis", s=50, alpha=0.75)
centers_orig = StandardScaler().fit(customers[features]).inverse_transform(km.cluster_centers_)
axes[0].scatter(centers_orig[:, 1], centers_orig[:, 2], marker="X", s=350,
                c="red", edgecolor="black", label="Centroids")
axes[0].set_xlabel("Annual Income"); axes[0].set_ylabel("Spending Score")
axes[0].set_title("Customer Segments"); axes[0].legend()

axes[1].scatter(customers["age"], customers["spending_score"],
                c=customers["segment"], cmap="viridis", s=50, alpha=0.75)
axes[1].set_xlabel("Age"); axes[1].set_ylabel("Spending Score")
axes[1].set_title("Age vs Spending")

plt.tight_layout(); plt.show()

K-Means++ Initialisation

Random initial centroids can produce poor local optima. k-means++ (scikit-learn's default, init="k-means++") spreads the initial centroids out: the first is random, and each subsequent one is chosen with probability proportional to its squared distance from the nearest existing centroid. This dramatically improves both quality and convergence speed.

Also set n_init=10 so the algorithm runs 10 times from different starts and keeps the best result by inertia.

Advantages and Disadvantages

K-Means vs K-Medoids

Applications

Customer segmentation, document/topic grouping, image compression (colour quantisation to k colours), delivery-zone planning, sensor/IoT grouping, anomaly detection (points far from every centroid), and numerosity reduction in preprocessing.

The next lesson covers the alternative that does not require you to pick k in advance: hierarchical clustering.