Module 7 · Unsupervised Learning ⏱ 19 min

K-Means, Built by Hand

By the end of this lesson you will be able to:
  • Implement the k-means assignment step: each point to its nearest centre
  • Implement the update step: move each centre to the mean of its assigned points
  • Run the assign-update loop to convergence with a fixed random seed

Clustering's most famous algorithm is built from two moves, and you already know both of them. Assign: every point joins the cluster of the centre it sits nearest. Update: each centre moves to the middle of the points now assigned to it. Alternate those two moves and the centres settle into place — that is the whole of k-means.

The name is literal: you choose a number of clusters, k, and the algorithm finds the means (averages) that best serve them. Nothing more. Its popularity comes from how little there is to it — a distance, a loop, and a stopping rule — and from how fast that loop runs on large data. In this lesson you write all three pieces yourself, run them on real points, and watch the centres drift into position.

flowchart TD
  I["pick k starting centers"] --> A["assign each point<br/>to nearest center"]
  A --> U["move each center<br/>to the mean of its points"]
  U --> CH["assignments changed?"]
  CH -->|"yes"| A
  CH -->|"no"| S["done"]
  style A fill:#10b981,color:#fff
  style U fill:#3776ab,color:#fff
  style S fill:#166534,color:#fff
K-means is a two-step loop: assign every point to its nearest centre, then move each centre to its members' mean. Repeat until nothing moves.

Step one: assign

The assignment step asks one question of each point: which centre is closest? You have written this already — it is the nearest logic from the previous lesson, now applied to every point in turn. The output is a list of labels, one per point, naming the centre each point currently belongs to.

Nothing about the centres changes during this step; they stay where they are while the points file past and choose. Notice that the result depends entirely on where the centres happen to be standing. With different starting centres, the very same points would be labelled differently — which is why the starting positions matter, a point we come back to at the end of the lesson.

The assignment step: every point labelled by its nearest centre.
import math

def distance(p, q):
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(p, q)))

def assign(points, centers):
    labels = []
    for p in points:
        best, best_d = 0, distance(p, centers[0])
        for i in range(1, len(centers)):
            d = distance(p, centers[i])
            if d < best_d:
                best, best_d = i, d
        labels.append(best)
    return labels

points = [(0, 0), (1, 0), (0, 1), (10, 10), (11, 10), (10, 11)]
centers = [(0, 0), (10, 10)]
print('labels:', assign(points, centers))

Step two: update

Once every point has a label, each centre moves to the mean of its members — the coordinate-wise average. Sum the positions of all the points in a cluster and divide by how many there are, separately for each feature. That average is the single position that minimises the total squared distance to the members, which is exactly why the method is called k-means.

After the move the centre sits in the middle of its cluster rather than wherever it happened to start. Run the loop and watch: a centre that began on top of a stray point slides toward the dense part of its group, assignment and update chasing each other until nothing moves any more.

One full iteration: assign, then move each centre to its members' mean.
import math

def distance(p, q):
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(p, q)))

def assign(points, centers):
    labels = []
    for p in points:
        best, best_d = 0, distance(p, centers[0])
        for i in range(1, len(centers)):
            d = distance(p, centers[i])
            if d < best_d:
                best, best_d = i, d
        labels.append(best)
    return labels

def update(points, labels, k):
    new_centers = []
    for c in range(k):
        members = [points[i] for i in range(len(points)) if labels[i] == c]
        centroid = tuple(sum(m[d] for m in members) / len(members) for d in range(len(members[0])))
        new_centers.append(centroid)
    return new_centers

points = [(0, 0), (1, 0), (0, 1), (10, 10), (11, 10), (10, 11)]
centers = [(0, 0), (10, 10)]
labels = assign(points, centers)
print('after assign:', labels)
centers = update(points, labels, 2)
print('after update :', [(round(c[0], 2), round(c[1], 2)) for c in centers])
flowchart LR
  subgraph I1["iteration 1"]
    P1["centers are guesses"] --> Q1["points reassigned"]
  end
  subgraph I2["iteration 2"]
    P2["centers drift toward<br/>their groups"] --> Q2["fewer reassignments"]
  end
  subgraph I3["converged"]
    P3["centers sit at group means"] --> Q3["no point switches"]
  end
  style P2 fill:#3776ab,color:#fff
  style P3 fill:#166534,color:#fff
Across iterations the centres drift toward their groups and fewer points switch, until the assignment is stable.

Convergence, and why the seed matters

Repeat assign-then-update until the centres stop moving. That fixed point is called convergence — the labels have stabilised, so the next update would reproduce the same centres, and the algorithm has nothing left to do. For ordinary, well-separated data it arrives in a handful of iterations.

The starting centres are chosen at random from the data points, which means k-means is not deterministic unless you pin that randomness down. Always pass a seed. Reproducible results are not a nicety here: without a fixed seed the same data can cluster differently on every run, and a result you cannot reproduce is one you cannot debug or defend. The discipline you learned tuning gradient descent carries over directly — control the randomness, or it controls you.

The full loop to convergence, seeded so the result is reproducible. Press Run.
import math, random

def distance(p, q):
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(p, q)))

def assign(points, centers):
    labels = []
    for p in points:
        best, best_d = 0, distance(p, centers[0])
        for i in range(1, len(centers)):
            d = distance(p, centers[i])
            if d < best_d:
                best, best_d = i, d
        labels.append(best)
    return labels

def update(points, labels, k):
    new_centers = []
    for c in range(k):
        members = [points[i] for i in range(len(points)) if labels[i] == c]
        centroid = tuple(sum(m[d] for m in members) / len(members) for d in range(len(members[0])))
        new_centers.append(centroid)
    return new_centers

def kmeans(points, k, seed=0, max_iters=100):
    rng = random.Random(seed)
    centers = rng.sample(list(points), k)
    for _ in range(max_iters):
        labels = assign(points, centers)
        new_centers = update(points, labels, k)
        if new_centers == centers:
            break
        centers = new_centers
    labels = assign(points, centers)
    return centers, labels

points = [(0, 0), (1, 0), (0, 1), (10, 10), (11, 10), (10, 11)]
centers, labels = kmeans(points, 2, seed=0)
print('labels :', labels)
print('centers:', [(round(c[0], 2), round(c[1], 2)) for c in centers])
flowchart TD
  G["the goal"] --> M["minimise total distance<br/>from points to their center"]
  M --> N["this quantity is inertia,<br/>the within-cluster sum of squares"]
  N --> H["each update lowers it<br/>until it settles"]
  style M fill:#10b981,color:#fff
Every update lowers (or holds) the total distance from points to their centres, called inertia — that monotonic descent is why k-means converges.
Exercise

Write assign(points, centers) returning a list of cluster labels — one per point — where each point takes the index of its nearest centre. Include distance so the function is self-contained.

import math

def distance(p, q):
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(p, q)))

def assign(points, centers):
    # one label per point: index of its nearest center
    pass
Exercise

What does this print? The centroid (mean) of three points along the diagonal.

def centroid(points):
    n = len(points)
    return tuple(sum(p[d] for p in points) / n for d in range(len(points[0])))

print(centroid([(0, 0), (2, 2), (4, 4)]))
Exercise

This update forgets to divide by the number of members, so it returns the SUM of each cluster's coordinates instead of the MEAN. Centres would shoot far away from their points. Add the missing division.

def update(points, labels, k):
    new_centers = []
    for c in range(k):
        members = [points[i] for i in range(len(points)) if labels[i] == c]
        # bug: this is the SUM, not the mean
        centroid = tuple(sum(m[d] for m in members) for d in range(len(members[0])))
        new_centers.append(centroid)
    return new_centers
Exercise

Write kmeans(points, k, seed=0, max_iters=100) using the provided assign and update. Pick k initial centres from points with random.Random(seed).sample, then loop assign-update until the centres stop changing, and return (centers, labels).

import math, random

def distance(p, q):
    return math.sqrt(sum((a - b) ** 2 for a, b in zip(p, q)))

def assign(points, centers):
    labels = []
    for p in points:
        best, best_d = 0, distance(p, centers[0])
        for i in range(1, len(centers)):
            d = distance(p, centers[i])
            if d < best_d:
                best, best_d = i, d
        labels.append(best)
    return labels

def update(points, labels, k):
    new_centers = []
    for c in range(k):
        members = [points[i] for i in range(len(points)) if labels[i] == c]
        centroid = tuple(sum(m[d] for m in members) / len(members) for d in range(len(members[0])))
        new_centers.append(centroid)
    return new_centers

def kmeans(points, k, seed=0, max_iters=100):
    # pick k starting centers, then assign-update until centers stop moving
    pass
Exercise

Write inertia(points, centers, labels) — the total squared distance from every point to its assigned centre. For each point, look up centers[labels[i]] and sum the squared differences. This is the quantity k-means is trying to minimise, and you will need it next lesson to choose k.

def inertia(points, centers, labels):
    # sum over points of squared distance to centers[labels[i]]
    pass

Recap

  • K-means alternates two moves: assign each point to its nearest centre, then update each centre to the mean of its members.
  • A centre's mean is the position that minimises total squared distance to its members — hence the name k-means.
  • The loop converges when the centres stop moving; for clean data that takes only a few iterations.
  • Always pass a seed so the random starting centres — and therefore your result — are reproducible.
  • Watch for empty clusters (a centre with no members) and local minima (a poor start trapped in a mediocre clustering); rerunning from several seeds is the standard defence.

You can now cluster a dataset, but k-means needs you to tell it how many clusters to find. The next lesson confronts that question: how do you choose k when nobody hands you the answer?

Checkpoint quiz

The k-means update step moves each centre to the ____ of the points assigned to it.

K-means is said to converge when:

Go deeper — technical resources