Clustering with k-means

So far every example came with a label. Clustering finds groups in data that has no labels: customers with similar habits, news stories on the same topic, pixels of the same colour.

The algorithm

k-means splits the points into k clusters, each with a centre, its centroid:

1. Start: pick k points as the first centroids.
2. Assign: put every point in the cluster of its nearest centroid.
3. Update: move each centroid to the mean of its cluster's points.
4. Repeat steps 2 and 3 until no point changes cluster.

Step through: the crosses are the centroids, and the faint lines show which centroid each point belongs to. In a couple of rounds the centroids slide into the middle of the three groups and stop.

points = [1, 2, 3, 10, 11, 12]
centroids = [1, 3]                       # start: two of the points

for round_number in range(1, 10):
    # Assign each point to its nearest centroid
    labels = [0 if abs(p - centroids[0]) <= abs(p - centroids[1]) else 1 for p in points]
    # Move each centroid to the mean of its points
    new = []
    for k in range(2):
        mine = [p for p, label in zip(points, labels) if label == k]
        new.append(sum(mine) / len(mine))
    print(f"round {round_number}: labels {labels}, centroids {new}")
    if new == centroids:
        break
    centroids = new

The program runs the same two steps on six points in a line. In round 1 the point 3 joins the right-hand cluster, because the starting centroid 3 is itself; once the centroids move, it switches sides. When the centroids stop moving, the program stops.

What k-means makes smaller

Each step lowers (or keeps) the total squared distance from the points to their centroids, the inertia. It can never go up, so k-means always stops. But it may stop in a local optimum: press Try other starting centroids in the simulator a few times, especially with k = 4 or 5, and the clusters can come out differently. Libraries therefore run k-means several times and keep the best result, and start with k-means++, which spreads the first centroids out.

Choosing k

More clusters always give a smaller inertia (with k = number of points, it is 0), so the smallest inertia is no guide. The elbow method plots inertia against k: it falls steeply, then flattens; the "elbow" where it bends is a good k. Often the right k comes from the problem itself: "we can run 4 marketing campaigns".

k-means finds round clusters of similar size. For long, curved or uneven clusters there are other methods, such as DBSCAN, which groups points that have many close neighbours.

Mark lesson complete