Classic ML · Easy · 15 min

k-means by hand

Assign, recompute, repeat: k-means is two functions and a loop. Cluster 90 customers by ticket and frequency and see how the elbow picks k.

POINTS are 90 customers as [mean ticket in EUR, visits per month]. Write five functions:

standardize(points): each column to (x - mean) / deviation (deviation over n); zeros when the deviation is 0.

assign(points, centroids): for each point, the index of the nearest centroid (squared distance, dist2); on a tie, the lower index.

update(points, labels, k, prev): k new centroids, each the mean of its points; a group with no points keeps prev[c].

kmeans(points, k, seed, maxIters): initial centroids = k distinct points picked with mulberry32(seed) (Math.floor(rng() * n), drawing again when the index already came out). Repeat assign and update up to maxIters times or until the labels stop changing. Return { centroids, labels, iterations, inertia }, with inertia the sum of dist2 from each point to its centroid.

elbow(points, ks, seed): the inertia of kmeans for each k in ks (with maxIters 50). The preview draws the groups, the centroids and the elbow curve.

Challenges 0/5

  • Standardise by columns, with zeros when nothing varies
  • Assign to the nearest centroid, lower index on a tie
  • Recompute centroids, keeping the one of an empty group
  • k-means with k = 3 on the standardised data separates the three groups
  • The elbow curve falls and flattens at k = 3

function standardize(points) {
  // per column: (x - mean) / sd with sd over n; zeros when sd is 0
  return points;
}

function assign(points, centroids) {
  // index of the nearest centroid by dist2; the lower index on a tie
  return points.map(() => 0);
}

function update(points, labels, k, prev) {
  // mean of the points of each group; prev[c] when the group is empty
  return prev;
}

function kmeans(points, k, seed, maxIters) {
  // init: k distinct points via mulberry32(seed); repeat assign + update until labels stop changing or maxIters
  // { centroids, labels, iterations, inertia }
  return { centroids: [], labels: points.map(() => 0), iterations: 0, inertia: 0 };
}

function elbow(points, ks, seed) {
  // kmeans(points, k, seed, 50).inertia for each k
  return ks.map(() => 0);
}

console.log(kmeans(standardize(POINTS), 3, 7, 50));
Console output appears here (console.log).

The idea

k-means knows nothing about customers. It starts from k random centres, assigns each point to the nearest centre, moves each centre to the mean of its points and repeats until nothing changes. It is two functions and a loop, and with 90 points it converges in under ten rounds.

The trap is the units. The mean ticket runs from 20 to 400 EUR and visits from 0.2 to 6: without standardising, the ticket decides the distance and visits do not count. Try the preview with and without standardising and you will see the groups change. That decision is yours, not the algorithm's.

Nor does the algorithm decide k. The elbow curve draws the inertia (sum of squared distances) for each k: it always falls, but it stops falling fast at the "natural" k. Here the elbow is at 3, which is how we built the data; on real data the elbow is usually less clear and you should look at the groups, not only the curve.

The seed fixes the result. With another seed, another start and sometimes another local minimum; scikit-learn solves that by starting ten times and keeping the best inertia. Here you do it once to see what is underneath.

Run it on your machine

With Python and scikit-learn, on this exercise's CSV (CC BY 4.0):

curl -O https://vorluxai.com/data/labs/clientes-kmeans.csv
python3 -c "
import pandas as pd
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
X = StandardScaler().fit_transform(pd.read_csv('clientes-kmeans.csv')[['gasto_medio_eur', 'visitas_mes']])
for k in range(1, 7):
    print(k, round(KMeans(n_clusters=k, n_init=10, random_state=0).fit(X).inertia_, 1))
"

Check what you learned

1. Why standardise before k-means on ticket and visits?
2. What does the elbow curve say?
3. Two seeds give different clusterings. What is going on?
4. What does `update` do with a group left without points?

Go deeper: the scikit-learn reference →

This in production, with your data? Let's talk for 15 minutes →