K-Means Clustering

L2-02 through L2-05 were supervised — labeled data, predict the label. K-Means is unsupervised — no labels, just find structure.

It’s the simplest practical unsupervised algorithm. Worth understanding deeply.

The Problem

Given a bunch of points, group them into k clusters so that:

• Points in the same cluster are similar (close)
• Points in different clusters are different (far)

Before:           After K-Means (k=3):
   *                  *  ●●
  * *                * * ●●
 * *   *            * *
   *  * *  *    →     * *   ▲
      *                     ▲▲
   *  *  *               *  ▲
   * *                  * * ▲

The Algorithm

K-Means iterates two steps until convergence:

Step 1: Assignment

Each point joins the nearest centroid.

Step 2: Update

Each centroid moves to the mean of its assigned points.

Repeat until centroids stop moving.

In Pseudocode

def kmeans(points, k):
    # Initialize: k random centroids
    centroids = random_sample(points, k)

    for _ in range(max_iter):
        # Step 1: assign each point to nearest centroid
        clusters = [[] for _ in range(k)]
        for p in points:
            nearest = argmin([distance(p, c) for c in centroids])
            clusters[nearest].append(p)

        # Step 2: move each centroid to mean of its cluster
        new_centroids = [mean(cluster) for cluster in clusters]

        if centroids == new_centroids:
            break
        centroids = new_centroids

    return centroids, clusters

That’s K-Means. ~10 lines.

Why It Works

Each step monotonically decreases the within-cluster sum of squares (WCSS):

$$\text{WCSS} = \sum_{i=1}^k \sum_{p \in C_i} \|p - \mu_i\|^2$$

• Assignment step: minimizes WCSS by moving each point to its nearest centroid
• Update step: minimizes WCSS by moving each centroid to the mean

WCSS is bounded below by 0, so iterations converge.

In Scikit-Learn

from sklearn.cluster import KMeans
import numpy as np

# Some data
X = np.random.randn(100, 2)

# Fit
kmeans = KMeans(n_clusters=3, random_state=42)
kmeans.fit(X)

# Get assignments
labels = kmeans.labels_         # [0, 1, 2, 0, ...]
centroids = kmeans.cluster_centers_

# Predict cluster of new point
new_points = np.array([[0.5, 0.5]])
cluster = kmeans.predict(new_points)

Choosing k

The biggest challenge: what’s the right k?

The Elbow Method

Plot WCSS for different k. Look for an “elbow” — where adding more clusters stops helping much:

wcss = []
for k in range(1, 11):
    km = KMeans(n_clusters=k).fit(X)
    wcss.append(km.inertia_)

import matplotlib.pyplot as plt
plt.plot(range(1, 11), wcss, marker='o')

If the plot shows a clear bend at k=3, use k=3.

Silhouette Score

A more rigorous metric:

from sklearn.metrics import silhouette_score

scores = []
for k in range(2, 11):
    km = KMeans(n_clusters=k).fit(X)
    score = silhouette_score(X, km.labels_)
    scores.append(score)
# Pick k with highest silhouette

Higher = better-separated clusters. Range [-1, 1].

Domain Knowledge

Sometimes you know:

• “We’re segmenting customers — 3-5 segments makes business sense”
• “Cluster news articles into topics — try 10-20”

Don’t rely solely on metrics. Use your understanding.

Initialization Matters

K-Means is sensitive to initial centroids. Random init can lead to local optima.

K-Means++ initialization (used by default in scikit-learn):

1. Pick first centroid randomly from data points
2. Pick next centroid with probability proportional to squared distance from existing centroids
3. Repeat until k centroids chosen

Spreads initial centroids → better convergence.

KMeans(n_clusters=5, init='k-means++', n_init=10)

n_init=10 runs the algorithm 10 times with different random inits, keeps the best — even more robust.

Limitations

1. Assumes spherical clusters

K-Means uses Euclidean distance → assumes clusters are round. If your data has crescent-shaped clusters, K-Means fails.

  K-Means thinks:        Real:
       ●●●                       ●●●
      ●●●●●         vs          ●●● ●●●
       ●●●                    ●●●     ●●●
                              ●●●     ●●●
                               ●●● ●●●
                                  ●●●

For non-spherical: use DBSCAN or HDBSCAN.

2. Needs k in advance

Unlike DBSCAN, you must specify k. Adds a hyperparameter.

3. Sensitive to outliers

A few outliers can pull centroids far from the cluster center.

Solution: remove outliers first, or use K-Medians.

4. Curse of dimensionality

In high-dim space, all points become roughly equidistant. K-Means struggles in 100D+.

Solution: reduce dim first with PCA / UMAP, then cluster.

Variants

Variant	Idea
K-Medians	Use median instead of mean — robust to outliers
K-Modes	For categorical data
K-Means++	Smart initialization (default in sklearn)
Mini-Batch K-Means	Faster for huge datasets
Bisecting K-Means	Hierarchical variant

When to Use K-Means

Good for:

• Customer segmentation (e.g., RFM)
• Image color quantization (reduce palette)
• Document grouping (after embedding)
• Anomaly detection (find points far from any cluster)
• Pre-processing for other models (cluster features as new feature)

Not good for:

• Non-spherical clusters
• Unknown k
• High-dim raw data (reduce dim first)
• Clusters of very different sizes

Real Example: Customer Segmentation

Group customers by spending behavior:

import pandas as pd
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler

# Features
df = pd.DataFrame({
    "annual_spend": [...],
    "purchase_frequency": [...],
    "recency_days": [...],
})

# Normalize (K-Means is scale-sensitive!)
scaler = StandardScaler()
X = scaler.fit_transform(df)

# Cluster
kmeans = KMeans(n_clusters=4, random_state=42)
df["segment"] = kmeans.fit_predict(X)

# Inspect segments
df.groupby("segment").mean()

Result: 4 segments with distinct profiles. Each segment can have its own marketing strategy.

A Visualization

Watch K-Means converge step-by-step on a simple dataset:

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

# Make blobs
from sklearn.datasets import make_blobs
X, _ = make_blobs(n_samples=300, centers=4, cluster_std=0.6)

# Plot iterations
fig, axes = plt.subplots(1, 4, figsize=(15, 3))
for i, ax in enumerate(axes):
    km = KMeans(n_clusters=4, n_init=1, max_iter=i+1)
    km.fit(X)
    ax.scatter(X[:, 0], X[:, 1], c=km.labels_, s=20)
    ax.scatter(km.cluster_centers_[:, 0], km.cluster_centers_[:, 1],
              c='red', s=200, marker='*')
    ax.set_title(f"After {i+1} iterations")
plt.show()

You see the centroids “pull” toward their final positions.

Always Normalize First

K-Means uses Euclidean distance — features with large ranges dominate.

# BAD: age (0-100) and salary (0-100k) — salary dominates
X = df[["age", "salary"]].values
kmeans.fit(X)

# GOOD: normalize first
from sklearn.preprocessing import StandardScaler
X = StandardScaler().fit_transform(df[["age", "salary"]].values)
kmeans.fit(X)

Forgetting to normalize is the most common K-Means bug.

Next recommended: L2-07 Evaluation Metrics or L2-04 Decision Trees.