k-means clustering

The concept

Clustering is unsupervised learning: the data have no labels, and the goal is to find groups of similar points. k-means, in the form known as Lloyd's algorithm, is the standard method. You choose the number of clusters k and some initial centroids, and then repeat two steps. In the assignment step, every point joins the cluster of its nearest centroid. In the update step, each centroid moves to the mean of the points assigned to it. The algorithm stops when the centroids no longer change.

Each step can only lower, or keep the same, the total squared distance of points to their centroids, so the algorithm always converges. It converges to a local optimum, though, and the result depends on the starting centroids. Practical code runs several starts and keeps the best, and k-means++ chooses initial centroids that are spread out. The number of clusters k must be chosen by you, for example with the elbow method or silhouette scores.

The example clusters seven 2-D points with k = 2, starting from the first and fourth points. It converges in 3 iterations to centroids (1.25, 1.5) and (3.9, 5.1), with 2 and 5 points. k-means works best for compact, roughly round clusters of similar size, and it needs numeric features that are on comparable scales.

Code example

points = [(1.0, 1.0), (1.5, 2.0), (3.0, 4.0), (5.0, 7.0), (3.5, 5.0), (4.5, 5.0), (3.5, 4.5)]


def distance2(a, b):
    return (a[0] - b[0]) ** 2 + (a[1] - b[1]) ** 2


def kmeans(points, centroids):
    iterations = 0
    while True:
        iterations += 1
        clusters = [[] for _ in centroids]
        for point in points:
            nearest = min(range(len(centroids)), key=lambda i: distance2(point, centroids[i]))
            clusters[nearest].append(point)
        updated = [(sum(p[0] for p in c) / len(c), sum(p[1] for p in c) / len(c)) for c in clusters]
        if updated == centroids:
            return centroids, clusters, iterations
        centroids = updated


centroids, clusters, iterations = kmeans(points, [points[0], points[3]])
print([tuple(round(value, 2) for value in c) for c in centroids])
print([len(c) for c in clusters], iterations)

[(1.25, 1.5), (3.9, 5.1)]
[2, 5] 3

Follow the code

Open the walkthrough to inspect the example one step at a time.

—

—

—

[(1.25, 1.5), (3.9, 5.1)]

[(1.25, 1.5), (3.9, 5.1)]
[2, 5] 3

Behind the scenes

Cost and common mistakes

Each iteration costs O(n·k·d) for n points, k clusters, and d dimensions. The number of iterations is usually small, but it has no tight general bound.

• k-means can converge to a poor local optimum. Run it from several initial centroids and keep the best clustering.
• A cluster can become empty, which would divide by zero when computing its mean. Real implementations re-seed empty clusters.
• k is not learned automatically. The data do not tell you the right number of clusters without extra analysis.

Try it yourself

Which points end up in the first cluster?

Further reading

• scikit-learn user guide (external)
• An Introduction to Statistical Learning (James, Witten, Hastie, Tibshirani) (external)
• Arthur and Vassilvitskii (2007): k-means++ (external)