K-Means visualized — centroid movement + elbow plot (iris)

SEE the centroids move iteration by iteration. SEE the elbow plot to pick K. Unsupervised learning made visual.

script.py v3.1

import numpy as np
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler

# ── Load iris (but pretend we DON'T know the labels — unsupervised!) ──
iris = load_iris()
X = iris.data[:, [2, 3]]  # petal length, petal width
# NOTE: we DON'T use y (labels) — K-Means is UNSUPERVISED
X_scaled = StandardScaler().fit_transform(X)

print(f'Dataset: iris (UNSUPERVISED — we pretend we don\'t know the species)')
print(f'Features: petal length, petal width (scaled)')
print(f'True number of species: 3 (but K-Means doesn\'t know this)')

# ════════════════════════════════════════════════════════════════════
#  PLOT 1: See centroids move over iterations (K=3)
# ════════════════════════════════════════════════════════════════════
fig, axes = plt.subplots(1, 4, figsize=(24, 5))

# Run K-Means with max_iter=1 to capture each step
kmeans_final = KMeans(n_clusters=3, random_state=42, n_init=10)
kmeans_final.fit(X_scaled)
final_centroids = kmeans_final.cluster_centers_

# Simulate iterations by running with different max_iter
for ax, max_iter in zip(axes, [1, 2, 5, 10]):
    km = KMeans(n_clusters=3, max_iter=max_iter, n_init=1, random_state=42, init='random')
    km.fit(X_scaled)

    # Plot data points colored by cluster assignment
    labels = km.labels_
    colors = ['#e74c3c', '#2ecc71', '#3498db']
    for i in range(3):
        mask = labels == i
        ax.scatter(X_scaled[mask, 0], X_scaled[mask, 1], c=colors[i], s=25, alpha=0.6)

    # Plot centroids (big X)
    ax.scatter(km.cluster_centers_[:, 0], km.cluster_centers_[:, 1],
              marker='X', s=200, c='black', edgecolors='yellow', linewidths=2,
              zorder=5, label='centroids')

    ax.set_title(f'Iteration {max_iter}\ninertia={km.inertia_:.1f}')
    ax.set_xlabel('petal length (scaled)')
    ax.set_ylabel('petal width (scaled)')
    ax.legend(fontsize=8)

plt.suptitle('K-Means: centroids move to minimize inertia', fontsize=14, y=1.02)
plt.tight_layout()

# ════════════════════════════════════════════════════════════════════
#  PLOT 2: Elbow plot — how to pick K
# ════════════════════════════════════════════════════════════════════
fig2, ax2 = plt.subplots(figsize=(10, 5))
k_values = range(1, 11)
inertias = []
for k in k_values:
    km = KMeans(n_clusters=k, random_state=42, n_init=10)
    km.fit(X_scaled)
    inertias.append(km.inertia_)

ax2.plot(k_values, inertias, 'o-', color='#a0c000', markersize=8, linewidth=2)
ax2.axvline(3, color='red', linestyle='--', alpha=0.5, label='Elbow at K=3 (true number of species)')
ax2.set_xlabel('K (number of clusters)')
ax2.set_ylabel('Inertia (sum of squared distances to centroid)')
ax2.set_title('Elbow Plot: pick K where the curve bends')
ax2.legend()
ax2.grid(True, alpha=0.3)
plt.tight_layout()

print(f'\nElbow plot (inertia for each K):')
for k, inertia in zip(k_values, inertias):
    marker = ' ← ELBOW (true K)' if k == 3 else ''
    print(f'  K={k:2d}: inertia={inertia:8.1f}{marker}')

print(f'\nTHE MATHS:')
print(f'  K-Means minimizes INERTIA (within-cluster sum of squares):')
print(f'  Inertia = Σ Σ ||x - centroid||²')
print(f'           (sum over all clusters, all points in each cluster)')
print(f'\n  Algorithm (Lloyd\'s algorithm):')
print(f'    1. Place K centroids randomly')
print(f'    2. ASSIGN: each point joins its nearest centroid')
print(f'    3. UPDATE: move each centroid to the MEAN of its points')
print(f'    4. Repeat 2-3 until centroids stop moving (convergence)')
print(f'\n  Inertia ALWAYS decreases with more K (more clusters = smaller distances)')
print(f'  The ELBOW is where adding more K gives diminishing returns')
print(f'  Look for the "bend" in the curve — that\'s the natural number of clusters')
print(f'\n  WARNING: K-Means assumes:')
print(f'    1. Clusters are SPHERICAL (fails on elongated shapes)')
print(f'    2. Clusters are SIMILAR SIZE (fails on very uneven sizes)')
print(f'    3. K is known (use elbow plot or silhouette score to pick)')
print(f'    4. Scale matters (always scale features first!)')

# Compare with true labels
print(f'\nK-Means clusters vs true species (adjusted Rand index):')
from sklearn.metrics import adjusted_rand_score
ari = adjusted_rand_score(iris.target, kmeans_final.labels_)
print(f'  ARI = {ari:.3f} (1=perfect, 0=random)')
print(f'  K-Means found the 3 species WITHOUT being told they exist!')

Python output click Run to execute

Click "Run" to execute the code on the left.
The output will appear here.

🧠 How it works

K-Means Clustering — VISUALIZED: UNSUPERVISED: no labels. The algorithm finds structure on its own. ALGORITHM (Lloyd's): (1) Place K centroids randomly. (2) ASSIGN: each point joins its nearest centroid. (3) UPDATE: move each centroid to the MEAN of its points. (4) Repeat until centroids stop moving. INERTIA (cost function): sum of squared distances from each point to its centroid. Inertia = Σ Σ ||x - centroid||². K-Means minimizes this. ELBOW PLOT: inertia vs K. Always decreasing (more K = smaller distances). Look for the 'bend' — that's the natural number of clusters. On iris: elbow at K=3. K-Means finds the 3 species WITHOUT BEING TOLD! CENTROID MOVEMENT: Iteration 1: random centroids → bad clustering. Iteration 2: centroids move to cluster means → better. Iteration 5+: centroids stabilize → convergence. WARNINGS: (1) Clusters are SPHERICAL — fails on elongated/irregular shapes (use DBSCAN instead). (2) Clusters are SIMILAR SIZE — fails on very uneven sizes. (3) K must be chosen — elbow plot or silhouette score. (4) ALWAYS scale first — large features dominate distance. (5) Random initialization — run multiple times (n_init=10) and pick the best. ADJUSTED RAND INDEX (ARI): compares K-Means clusters to true labels. 1.0 = perfect match. 0.0 = random. On iris: ~0.7-0.9 — K-Means recovers the species!

✏️ Try these changes

Change n_clusters from 3 to 2
→ merges 2 species into 1 cluster — look at the elbow plot to see why 3 is better
Change n_clusters from 3 to 5
→ splits species into sub-clusters — overfitting the clustering
Change init from 'random' to 'k-means++'
→ smarter initialization — converges faster, better results
Use sepal features: X = iris.data[:, [0, 1]]
→ clusters are less clear — species overlap in sepal dimensions

← Previous lab

KNN visualized — the K nearest neighbors + decision boundary (iris)

Next lab →

Gradient Boosting visualized — sequential learning + residuals (breast cancer)