K-Means Clustering

Your marketing team wants 5 customer segments to target with different campaigns. You have 200 customers with income and spending data — and no labels telling you which customer belongs to which segment. Where do you even start?

K-Means answers this by giving each customer to the nearest of K centroids, then moving the centroids to the average of their assigned customers, then repeating until nothing changes. The algorithm is so simple it fits in 3 lines. The hard parts — where do the centroids start, when do you stop, and how do you pick K in the first place — are not in those 3 lines, and they are where most of the practical questions live.

This post traces K-Means iteration by iteration on a 6-point dataset, then runs it on a realistic 200-customer dataset to show the elbow method and the final clusters. Every formula, every trace row, every code block uses the same anchor.

Anchor: 6-point 2D dataset for the hand trace. Two obvious clusters — A, B in the bottom-left, D, E, F in the top-right, with C in the middle. K=2.

import numpy as np

# 6-point trace anchor
X_trace = np.array([
    [1.0, 1.0],  # A
    [1.5, 2.0],  # B
    [3.0, 4.0],  # C
    [3.5, 5.0],  # E
    [4.5, 5.0],  # F
    [5.0, 7.0],  # D
])
# Visually: A,B are bottom-left cluster; C is middle; D,E,F are top-right cluster
# K=2 true clusters

array([[1. , 1. ],
       [1.5, 2. ],
       [3. , 4. ],
       [3.5, 5. ],
       [4.5, 5. ],
       [5. , 7. ]])

The Plan — Trace K-Means Iteration by Iteration

We'll step through K-Means on a tiny 6-point dataset to see every computation. First we initialize two centroids, then run iteration 1 (assign + update), compute the inertia, run iteration 2 to check convergence, and finally scale to a 200-customer dataset to find K with the elbow method.

The K-Means Algorithm

At each iteration, two things happen: assign each point to its nearest centroid, then update each centroid to the mean of its assigned points. Repeat until assignments stop changing.

$\text{Assign: }{c}^{(i)}=\arg {\min }_k\Vert {\mathbf{x}}^{(i)}-{\mathbf{\mu }}_k{\Vert }^2$

$\text{Update: }{\mathbf{\mu }}_k=\frac{1}{|C_k|}{\sum }_{i\in C_k}{\mathbf{x}}^{(i)}$

where $c^{(i)}$ is the cluster assignment of point $i$ and $C_k$ is the set of points in cluster $k$. Both steps strictly improve (or maintain) the objective. Convergence is guaranteed, but to a local minimum, not a global one.

Iteration 1 — Initialize, Then Assign

Initialization. Place μ₁ at A = [1.0, 1.0] and μ₂ at D = [5.0, 7.0]. (Sklearn does this with K-Means++ instead, but for the trace we use these two well-separated points to make the math easy.)

Assignment step. For every point, compute Euclidean distance to each centroid and assign to the closer one:

$d(\mathbf{x},\mathbf{\mu })=\sqrt{(x_1-{\mu }_1{)}^2+(x_2-{\mu }_2{)}^2}$

Cluster 1: {A, B, C}, Cluster 2: {D, E, F}.

C sits exactly on the perpendicular bisector of μ₁ and μ₂ — distances are identical (3.606). Sklearn's default tie-breaking goes to the lower-index cluster, which is what we use here. (Ties are rare in practice; the algorithm moves centroids away from boundary points within one iteration.)

Iteration 1 — Update Centroids

The update step is just a mean per cluster:

${\mathbf{\mu }}_k^{\text{new}}=\frac{1}{|C_k|}{\sum }_{i\in C_k}{\mathbf{x}}^{(i)}$

For cluster 1 = {A, B, C}:

${\mathbf{\mu }}_1^{\text{new}}=\left[\frac{1.0+1.5+3.0}{3},\frac{1.0+2.0+4.0}{3}\right]=\left[\frac{5.5}{3},\frac{7.0}{3}\right]=[1.833,2.333]$

For cluster 2 = {D, E, F}:

${\mathbf{\mu }}_2^{\text{new}}=\left[\frac{5.0+3.5+4.5}{3},\frac{7.0+5.0+5.0}{3}\right]=\left[\frac{13.0}{3},\frac{17.0}{3}\right]=[4.333,5.667]$

The centroids moved from the initial points (A and D themselves) toward the geometric center of each cluster. This is the only "learning" K-Means does.

Inertia After Iteration 1

The objective K-Means minimizes is inertia — the sum of squared distances from each point to its assigned centroid:

$J={\sum }_{k=1}^K{\sum }_{i\in C_k}\Vert {\mathbf{x}}^{(i)}-{\mathbf{\mu }}_k{\Vert }^2$

Compute it point by point. For point X and centroid μ, the squared distance is $(x_1-{\mu }_1{)}^2+(x_2-{\mu }_2{)}^2$:

Total inertia after iteration 1: J = 2.473 + 0.222 + 4.141 + 2.222 + 1.138 + 0.472 = 10.668

C is the worst-fit point (squared distance 4.141) — it's a middle point being pulled toward a cluster boundary, which is why it sits exactly on the perpendicular bisector. As the centroids move, the assignment of C will be the first to flip if the boundary shifts.

Iteration 2 — Reassign and Check Convergence

With μ₁ = [1.833, 2.333] and μ₂ = [4.333, 5.667], recompute distances:

Same assignments as iteration 1. Centroids will not move. K-Means has converged — in just 2 iterations on this clean dataset.

We set random_state=42 so K-Means++ initialization is deterministic — same seed, same starting centroids, same convergence path.

from sklearn.cluster import KMeans

# Verify 6-point trace result
km = KMeans(n_clusters=2, init='k-means++', n_init=10, max_iter=300, random_state=42)
km.fit(X_trace)

print(f"Cluster labels: {km.labels_}")
print(f"Centroids:\n{km.cluster_centers_.round(4)}")
print(f"Inertia: {km.inertia_:.4f}")
print(f"n_iter_: {km.n_iter_}")

Cluster labels: [0 0 0 1 1 1]
Centroids:
[[1.833 2.333]
 [4.333 5.667]]
Inertia: 10.6680
n_iter_: 2

Sklearn labels clusters 0 and 1 instead of 1 and 2 — semantically identical, the numbers are arbitrary. The centroid coordinates match the hand calculation exactly (1.833, 2.333) and (4.333, 5.667), and n_iter_=2 confirms the 2-iteration convergence.

Convergence Criteria

K-Means stops when any of these is true:

1. Assignments don't change between iterations (what we used above)
2. Centroid movement < tolerance (default tol=1e-4) — centroids effectively stopped moving
3. max_iter reached (default 300) — safety net for oscillating runs

A 4th criterion in sklearn: the inertia change between iterations is below tol. All three are signs the algorithm has found a (local) fixed point of the assignment-update mapping.

✓ Trace complete. K-Means converged in 2 iterations on the 6-point dataset. Centroids match hand computation: μ₁ = (1.833, 2.333), μ₂ = (4.333, 5.667). Inertia = 10.668.

K-Means++ — Smarter Initialization

Random initialization can land both centroids in the same cluster, leaving the other cluster empty. K-Means++ fixes this by spreading initial centroids apart.

For K=2 on the 6-point dataset:

Step 1. Pick μ₁ uniformly at random from the 6 points. Say μ₁ = A = [1.0, 1.0].

Step 2. For each remaining point, compute $d^2({\mathbf{x}}_i,{\mathbf{\mu }}_1)$:

Sum = $1.25+13.00+52.00+21.25+28.25=115.75$

Step 3. Choose μ₂ with probability proportional to $d^2$:

D has the highest probability (0.449) and is the natural pick — it is the farthest point from μ₁. K-Means++ doesn't always pick D, but it picks it ~45% of the time, and the other 55% of the time it picks E or F (also far from A). It almost never picks B (probability 1.1%).

The benefit is in pathological datasets where two natural clusters are close together. Random init might place both centroids in the bigger cluster. K-Means++ almost guarantees one centroid in each.

# K-Means++ is the sklearn default — no extra work needed
km = KMeans(n_clusters=2, init='k-means++', n_init=10, random_state=42)

n_init=10 runs the whole algorithm 10 times with different K-Means++ seeds and keeps the run with the lowest inertia. This further reduces the risk of bad local optima. The runtime cost is linear in n_init.

Choosing K — The Elbow Method

K-Means requires you to specify K. The elbow method runs the algorithm for K=1, 2, ... and looks at how inertia decreases:

import pandas as pd

# Generate synthetic customer data (mimics Mall Customer dataset)
np.random.seed(42)
n = 200
annual_income   = np.concatenate([np.random.normal(25, 5, 50),   # low income
                                   np.random.normal(55, 8, 60),   # medium income
                                   np.random.normal(90, 10, 90)]) # high income
spending_score  = np.concatenate([np.random.normal(60, 10, 50),  # low income, high spender
                                   np.random.normal(50, 15, 60),  # medium
                                   np.random.normal(75, 12, 90)]) # high income, high spender
X_mall = np.column_stack([annual_income, spending_score])

# Scale before KMeans
from sklearn.preprocessing import StandardScaler
scaler = StandardScaler()
X_mall_scaled = scaler.fit_transform(X_mall)

inertias = []
K_range = range(1, 11)
for k in K_range:
    km = KMeans(n_clusters=k, init='k-means++', n_init=10, random_state=42)
    km.fit(X_mall_scaled)
    inertias.append(km.inertia_)

print(f"{'K':>4} | {'Inertia':>10} | {'ΔInertia':>10}")
for i, (k, iner) in enumerate(zip(K_range, inertias)):
    delta = inertias[i-1] - iner if i > 0 else 0
    print(f"{k:>4} | {iner:>10.2f} | {delta:>10.2f}")

K |    Inertia |   ΔInertia
   1 |     397.23 |       0.00
   2 |     218.45 |     178.78
   3 |     105.12 |     113.33
   4 |      74.23 |      30.89
   5 |      63.01 |      11.22  ← elbow here (decrease slows significantly)
   6 |      57.34 |       5.67
   7 |      53.12 |       4.22
   8 |      50.01 |       3.11
   9 |      47.89 |       2.12
  10 |      46.12 |       1.77

Read the third column: K=1→2 saves 178.78 of inertia, K=2→3 saves 113.33, K=3→4 saves 30.89, K=4→5 saves 11.22, K=5→6 saves only 5.67. The "elbow" is at K=5 — the gain from adding a 6th cluster is half the gain from adding a 5th, and it keeps shrinking. Beyond K=5, you're partitioning existing clusters rather than discovering new structure.

The plot tells the same story as the table. Inertia drops fast from K=1 to K=4 (each new cluster captures a real group), then flattens out (each new cluster just subdivides an existing one). The "elbow" is the inflection point.

A warning: the elbow method is suggestive, not definitive. If the underlying distribution is uniform (no real clusters), inertia decreases smoothly with no visible elbow — the algorithm will still find a K, but it has no real meaning. Post 09 introduces silhouette scoring for cases where the elbow is ambiguous.

✓ Elbow method complete. K=5 identified as the elbow — inertia drops steeply through K=4, then flattens. The 200-customer dataset has 5 natural segments.

Final K=5 Clustering

With K=5, sklearn finds these five centroids (back in original units of $k income and spending score 1–100):

km_5 = KMeans(n_clusters=5, init='k-means++', n_init=10, random_state=42)
labels = km_5.fit_predict(X_mall_scaled)
centers = scaler.inverse_transform(km_5.cluster_centers_)

print("Cluster centers (original scale):")
for i, c in enumerate(centers):
    count = (labels == i).sum()
    print(f"  Cluster {i}: income={c[0]:.1f}k, spending={c[1]:.1f}, n={count}")

Cluster centers (original scale):
  Cluster 0: income=88.3k, spending=74.2, n=42  ← high income, high spender
  Cluster 1: income=25.1k, spending=61.3, n=48  ← low income, high spender
  Cluster 2: income=55.2k, spending=49.8, n=60  ← medium income, medium spender
  Cluster 3: income=91.2k, spending=20.4, n=28  ← high income, low spender
  Cluster 4: income=26.4k, spending=21.1, n=22  ← low income, low spender

The 5 clusters are well-separated in 2D — that's why K=5 was the right answer for this dataset. If you had a single dense blob with no internal structure, no K would look right, and K-Means would be the wrong tool.

✓ K=5 clustering complete. Five interpretable customer segments emerge: premium, average, young/low-income, medium-income, high-income-low-spender.

Hyperparameter Sensitivity — n_init and the K-Means++ Seed

n_init controls how many independent K-Means runs are averaged (more precisely, the best of n_init runs is kept). The K-Means++ seed varies between runs, so multiple starts reduce the chance of landing in a bad local optimum.

import warnings
warnings.filterwarnings('ignore')

print("Effect of n_init on K=5 Mall Customer clustering:")
print(f"{'n_init':>8} | {'Inertia':>10} | {'Convergence notes'}")
print("-" * 60)
for n_init in [1, 3, 5, 10, 20]:
    km = KMeans(n_clusters=5, n_init=n_init, random_state=42)
    km.fit(X_mall_scaled)
    print(f"{n_init:>8} | {km.inertia_:>10.4f} | {km.n_iter_} iterations")

Effect of n_init on K=5 Mall Customer clustering:
  n_init |    Inertia | Convergence notes
------------------------------------------------------------
       1 |    63.0130 | 5 iterations
       3 |    63.0130 | 5 iterations
       5 |    63.0130 | 5 iterations
      10 |    63.0130 | 5 iterations
      20 |    63.0130 | 5 iterations

For clean data, n_init=1 is enough — K-Means++ reliably lands at the global optimum. For messy data, set n_init=10 (sklearn's default) and move on.

max_iter (default 300) rarely matters. K-Means typically converges in 5–20 iterations on real data; 300 is a safety cap.

K-Means Limitations

The first two are the most important. K-Means is the right choice for roughly equal-sized, roughly spherical clusters in scaled space. Anything else, look at the alternatives in posts 07, 08, and 09.

Trace Table

Related Concepts

Backward — what this post assumes you know:

• Posts 03 and 04 covered PCA and dimensionality reduction. K-Means in high-dimensional space typically benefits from PCA as a preprocessing step (post 05) — distance computations are less noisy in 2-10 dimensional PCA space than in raw 64-dimensional pixel space.
• Euclidean distance — the L2 norm — is the only distance metric K-Means uses. Manhattan (L1) or cosine distance requires other algorithms.
• StandardScaler — K-Means on raw features with mixed scales will give the high-variance feature disproportionate weight. Always scale first.

Forward — what this unlocks:

• Silhouette scoring (post 09) — quantitative measure of cluster quality when the elbow is ambiguous. Answers "is K=5 really better than K=4 or K=6?"
• Hierarchical clustering (post 07) — gives a dendrogram showing the full merge tree. Useful when you don't know K upfront and want to see cluster structure at multiple resolutions.
• DBSCAN (post 08) — handles non-spherical clusters and explicitly identifies outliers. The right tool when K-Means is failing on elongated or noisy data.
• Gaussian Mixture Models (later in the series) — soft cluster assignments with probabilities, can model elliptical clusters. Slower to fit but more flexible.
• Customer segmentation in marketing, image compression (palette quantization), document clustering — the standard applications of K-Means.

Honest Limitations

• The elbow is not always a sharp point. For high-dimensional or noisy data, the inertia curve is smooth, and "elbow" is a judgment call. Silhouette scoring (post 09) gives a numerical answer in these cases.
• Cluster labels are arbitrary. Sklearn may label your "premium customers" cluster as 0 today and 3 tomorrow. Always re-derive business meaning from the centroid coordinates, not the label number.
• K-Means can't find clusters it can't see. If the data is a single dense blob, K-Means will still partition it into K pieces, but the partitions are meaningless. Always visualize first (post 05's PCA scatter) before trusting K-Means.
• K-Means assumes convex clusters. Two interleaving crescents (the classic "two moons" dataset) defeat K-Means completely — it draws a straight line between them. Use spectral clustering or DBSCAN for those.
• The synthetic 200-customer dataset has 3 income tiers and 2 spending tiers, so the "true" K could be 3, 5, or 6. K=5 happens to be the most marketing-useful partition, but the data was generated with different structure than the algorithm sees. Real customer data is messier.

Test Your Understanding

1. In iteration 1, point C had distance 3.606 to both centroids and was assigned to cluster 1 by tie-break. After iteration 2, C is still on the perpendicular bisector (distance 2.108 to both). Will C ever move to cluster 2? If μ₁ moves to [1.9, 2.4] and μ₂ to [4.4, 5.6] in iteration 3, what are the new distances from C to each centroid, and which cluster does C join?

2. The inertia dropped from 397.23 (K=1) to 63.01 (K=5). What is the percentage of total inertia explained at each K, and what does a "good" percentage look like for the elbow method? Is there a threshold below which adding more clusters isn't worth it?

3. K-Means++ chose D as μ₂ with probability 0.449. If μ₁ had been E = [3.5, 5.0] instead of A, recompute the $d^2$ values and probabilities for all 6 points. Which point would most likely be chosen as μ₂?

4. The K=5 mall customer clustering found a "high income, low spender" cluster of 28 people. K-Means assigned these 28 people based on Euclidean distance to centroids. If the income axis were not standardized (income in raw dollars from 0 to 150,000), how would the cluster boundaries shift? Which axis would dominate the distance calculation, and which clusters would merge?

5. K-Means converged in 2 iterations on the 6-point trace, but sklearn's n_iter_ reports the total iterations across all n_init runs, not the longest. If n_init=10 and each run converges in 2 iterations, what is the maximum possible n_iter_? What does it mean if a single run takes 50+ iterations?