K-Means Clustering

What is clustering?

So far, many algorithms you’ve seen are given the right answers during training (that’s supervised learning). Clustering is different: we only have raw data and we ask the computer to find natural groups in it by itself. This is unsupervised learning.

Example: an online shop has data on thousands of customers — age, income and how much they spend. Nobody has labelled them. Clustering might discover groups like “young bargain-hunters”, “high-spending professionals” and “occasional buyers” — and the shop can target each group differently.

In the 3D model each dot is a customer and each axis is one feature. Points that are close together are similar.

The K-Means algorithm

You choose k, the number of groups you want. Then:

1. Initialise: pick k starting points called centroids (the ◆ shapes).
2. Assign step: every point joins the cluster of its nearest centroid (straight-line distance).
3. Update step: move every centroid to the mean (average position) of the points in its cluster.
4. Repeat steps 2 and 3 until no point changes cluster.

That’s it! In the model, the thin lines show which centroid each point belongs to. Watch the centroids glide towards the centre of their groups.

Measuring quality: WCSS

How good is a clustering? K-Means tries to minimise the Within-Cluster Sum of Squares (WCSS, also called inertia): the total squared distance from every point to its centroid.

WCSS = Σ over clusters Σ over points in cluster  ‖point − centroid‖²

Both steps can only lower WCSS (or keep it equal), so the algorithm is guaranteed to stop. The model shows WCSS dropping after every update.

Choosing k — the elbow method

K-Means can’t tell you how many groups there should be. A popular trick: run it for k = 1, 2, 3, … and plot WCSS. The curve drops sharply at first and then flattens. The “elbow” where it bends is usually a good k.

Problems and fixes

Code

From scratch with NumPy:

import numpy as np

def k_means(X, k, iterations=100):
    rng = np.random.default_rng()
    centroids = X[rng.choice(len(X), k, replace=False)]      # random start
    for _ in range(iterations):
        # assign: distance from every point to every centroid
        dists = np.linalg.norm(X[:, None, :] - centroids[None, :, :], axis=2)
        labels = dists.argmin(axis=1)
        # update: mean of each cluster (keep old centroid if a cluster is empty)
        new = np.array([X[labels == j].mean(axis=0) if np.any(labels == j) else centroids[j]
                        for j in range(k)])
        if np.allclose(new, centroids):
            break                                             # converged
        centroids = new
    return labels, centroids

X = np.vstack([np.random.randn(30, 3) + c for c in ([0, 0, 0], [6, 6, 0], [0, 6, 6])])
labels, centroids = k_means(X, k=3)
print(centroids.round(2))

In practice you would use scikit-learn, which uses k-means++ by default:

from sklearn.cluster import KMeans

model = KMeans(n_clusters=3, n_init=10).fit(X)
print(model.labels_[:10])        # cluster of the first 10 points
print(model.inertia_)            # WCSS

Where is K-Means used?

• Customer segmentation for marketing.
• Image compression: cluster the colours of an image into k colours.
• Document grouping: news articles about similar topics.
• Anomaly detection: points far from every centroid are suspicious.
• As a first step for other algorithms (e.g. initialising other models).

Common mistakes

• Forgetting to scale the features.
• Treating cluster numbers as meaningful: “cluster 2” is just a name; it can change between runs.
• Expecting K-Means to find long, curved or very unequal clusters — it prefers compact round ones.