Python k-means clustering from scratch

k-means is an unsupervised algorithm: there are no labels to predict. It repeatedly assigns each point to the nearest centroid, then moves each centroid to the mean of its members. The loop is short, but distance scaling and initialization decide what it finds. You will also try farthest-first initialization and compare k with inertia.

Measure distance

The usual choice is Euclidean distance. Squared distance preserves the same ordering and avoids a square root, but Euclidean distance is easier to read when explaining the result. Keep the coordinate order consistent.

import math

def distance(a, b):
    return math.dist(a, b)

print(distance((0, 0), (3, 4)))

Five units apart

Assign each point to a centroid

For every point, compare its distance to every centroid and keep the nearest index. A deterministic tie rule prevents the assignment from changing just because dictionary or set order changed.

def distance(a, b):
    return sum((x - y) ** 2 for x, y in zip(a, b)) ** 0.5

centroids = [(0, 0), (5, 5)]
point = (1, 1)
label = min(range(len(centroids)), key=lambda i: (distance(point, centroids[i]), i))
print(label)

Nearest centroid wins

Move each centroid to its mean

After assignment, each centroid becomes the coordinate-wise mean of the points labeled with it. An empty cluster has no members, so the code must decide whether to keep the old centroid, reseed it, or reduce k.

points = [(0, 0), (0, 1), (5, 5), (5, 6)]
labels = [0, 0, 1, 1]
members = [point for point, label in zip(points, labels) if label == 0]
centroid = (
    sum(x for x, _ in members) / len(members),
    sum(y for _, y in members) / len(members),
)
print(tuple(round(value, 1) for value in centroid))

The mean of the members

Iterate until labels stop changing

Repeat assignment and update until the labels or centroids stop moving, or until a step budget expires. The result is a local optimum, so a different initialization can produce a different clustering.

def distance(a, b):
    return sum((x - y) ** 2 for x, y in zip(a, b)) ** 0.5

points = [(0, 0), (0, 1), (5, 5), (5, 6)]
centroids = [(0, 0), (5, 5)]
labels = None
for _ in range(10):
    new_labels = [
        min(range(len(centroids)), key=lambda i: (distance(point, centroids[i]), i))
        for point in points
    ]
    new_centroids = []
    for index in range(len(centroids)):
        members = [point for point, label in zip(points, new_labels) if label == index]
        if members:
            new_centroids.append((
                sum(x for x, _ in members) / len(members),
                sum(y for _, y in members) / len(members),
            ))
        else:
            new_centroids.append(centroids[index])
    if new_labels == labels and new_centroids == centroids:
        break
    labels = new_labels
    centroids = new_centroids
print(labels)
print([(round(x, 1), round(y, 1)) for x, y in centroids])

Two clean groups

Farthest-first initialization spreads centroids

Random initialization can place two centroids in the same dense region. Farthest-first picks the first point, then repeatedly adds the point with the largest distance to its nearest centroid. It is deterministic given a start point and often gives a more useful initial spread.

def distance(a, b):
    return sum((x - y) ** 2 for x, y in zip(a, b)) ** 0.5

points = [(0, 0), (1, 0), (5, 5), (6, 5)]
centroids = [points[0]]
while len(centroids) < 2:
    candidate = max(points, key=lambda point: min(distance(point, c) for c in centroids))
    centroids.append(candidate)
print(centroids)

Choose the next centroid far away

Inertia helps compare k, but does not prove it

Inertia is the sum of squared distances from points to their assigned centroid. It usually falls as k grows, so the smallest value is not automatically the best choice. Use the elbow pattern, domain constraints, and stability across initializations, and handle empty clusters explicitly.

def inertia(points, centroids):
    total = 0
    for point in points:
        nearest = min(centroids, key=lambda c: sum((x - y) ** 2 for x, y in zip(point, c)))
        total += sum((x - y) ** 2 for x, y in zip(point, nearest))
    return total

points = [(0, 0), (1, 0), (5, 5), (6, 5)]
for k, centroids in [(1, [(3, 2.5)]), (2, [(0.5, 0), (5.5, 5)])]:
    print(k, round(inertia(points, centroids), 1))

Lower inertia is not an automatic winner

Try it yourself

Change the values, predict the output, then run the code. The Python runtime is downloaded the first time you press Run, so the first run can take a few seconds.