Inventing K-Means Clustering

Here's a bunch of points scattered on a 2D plane. Your goal: split them into two groups so that points in the same group are close together.

How would you start? You need some notion of a "center" for each group. The simplest idea: just pick two points and call them the centers. Then assign every point to whichever center is closer.

Before you can group anything, you need to measure how far apart two points are.

If you completed the Expressions & Functions chapter, you already wrote this in Exercise 3.3 — "Distance Between Two Points (2D)." You can reuse that function here.

Write distance(p1, p2) that computes the straight-line (Euclidean) distance between two 2D points.

Recall: the distance between $(x_1, y_1)$ and $(x_2, y_2)$ is $\sqrt{(x_1 - x_2)^2 + (y_1 - y_2)^2}$.

You have 10 points and 2 centers. For each point, figure out which center is closer, and put that point in that center's group.

Write assign_to_groups(points, centers) that returns a list of groups — one group per center. Each group is a list of points assigned to that center.

You have two groups. The centers you started with were arbitrary — just the first two points. Now that you have groups, can you find better centers?

What single point best represents a group of points?

Write compute_centers(groups) that returns a new list of centers — one per group.

You can now assign points to centers and compute better centers from the groups. What happens if you keep doing this — assign, recompute, assign, recompute?

Starting from your initial centers ([1, 2] and [8, 8]):

1. Assign all points to groups using assign_to_groups
2. Compute new centers using compute_centers
3. Print the old and new centers

Did the centers move?

Run the assign → recompute cycle multiple times. Print the centers after each round.

When should you stop? Look at the centers each round — what do you notice?

Write a loop that repeats the cycle and stops when the centers don't change (or change by less than a tiny amount like 0.0001).

Run the plotting helper below to visualize the groups your algorithm found. The centers are shown as large X markers.

Don't worry about understanding the plotting code — just run it and look at the result.

Congratulations! You just invented one of the most widely used algorithms in machine learning: K-Means clustering.

Here's what you built:

K-Means is an unsupervised learning algorithm — it finds structure in data without any labels. No one told the algorithm which points belong together. It figured it out on its own, using nothing but distances and averages.

Your code works for 2 groups in 2D. But what if you have 3D data? Or want 5 groups? Let's make it general.

Your distance function only works for 2D points. Rewrite it to work with points of any number of dimensions.

The Euclidean distance in N dimensions is: $\sqrt{(a_0 - b_0)^2 + (a_1 - b_1)^2 + \ldots + (a_{n-1} - b_{n-1})^2}$

Your compute_centers function only averages x and y. Rewrite it to work with points of any number of dimensions.

Now combine everything into a single function that works for any number of groups (K) and any number of dimensions.

Write:

def kmeans(points, k, max_iterations=100):
    """
    Args:
        points: list of points (each point is a list of numbers)
        k: number of groups to create
        max_iterations: stop after this many iterations even if not converged
    Returns:
        (centers, groups) — final centers and the list of groups
    """

Use the first K points as the initial centers.

Different problems call for different notions of "distance." For example, Manhattan distance (sum of absolute differences) is better when your coordinates represent things like city blocks.

Modify your kmeans function to accept a distance function as an argument:

def kmeans(points, k, distance_fn=distance, max_iterations=100):

Also update assign_to_groups to accept distance_fn.

Test it with both Euclidean and Manhattan distance.

K-Means needs you to choose K — the number of groups. But how do you know the right K? What if you pick too many? Too few?

Let's find out by measuring how "good" each grouping is.

How tight are your clusters? One measure: for each point, compute its distance to its center. Sum up all these distances. A lower total means tighter clusters.

Write compute_inertia(centers, groups, distance_fn=distance) that returns this total.

Run K-Means on the dataset below with k = 2, 3, 4, 5, and 6. Compute the inertia for each. What happens to the inertia as K increases?

Plot inertia vs. K. The curve always goes down, but at some point the drop slows dramatically — like the bend of an elbow. That bend is your best K.

Where is the elbow for the dataset above? Does it match the number of clusters we generated (3)?

See what the clusters look like for different values of K. Run the code below to plot k=2, k=3, and k=5 side by side.

Which one looks most natural?