Clustering Algorithms in Python

K-Means

Pick k starting centroids, assign every point to its nearest centroid, then recompute each centroid as the mean position of the points now assigned to it. Repeat assign-then-recompute until no point changes clusters (or, equivalently, no centroid moves) — that's the whole algorithm. It always converges to a local optimum, but which one depends heavily on where the centroids started.

k has to be chosen up front — K-Means has no built-in notion of "how many clusters are actually here," it will happily force any k you give it, even a bad one.

Worked Example

Seven 2D points, split visually into a tight group near the origin and a looser group up and to the right: P1(1,1), P2(1.5,2), P3(3,4), P4(5,7), P5(3.5,5), P6(4.5,5), P7(3.5,4.5). k=2, with deterministic (non-random) initial centroids at P1 and P2 so the trace below is exactly reproducible: C0=(1,1), C1=(1.5,2).

Iteration 1 moves C1 a long way (it starts equal to P2, then jumps to the mean of six points once everything but P1 lands in its cluster). By iteration 3 the assignments and centroids both stop changing, so the loop exits. Final clusters: {P1, P2} and {P3, P4, P5, P6, P7}.

import math

def euclidean_distance(a, b):
    return math.sqrt(sum((a[i] - b[i]) ** 2 for i in range(len(a))))

def kmeans(points, k, initial_centroids, max_iters=100):
    centroids = list(initial_centroids)
    for _ in range(max_iters):
        # assign each point to its nearest centroid
        assignments = []
        for p in points:
            distances = [euclidean_distance(p, c) for c in centroids]
            assignments.append(distances.index(min(distances)))

        # recompute each centroid as the mean of its assigned points
        new_centroids = []
        for cluster_id in range(k):
            members = [points[i] for i in range(len(points)) if assignments[i] == cluster_id]
            if not members:
                # empty cluster -- keep the old centroid rather than dividing by zero
                new_centroids.append(centroids[cluster_id])
                continue
            mean_x = sum(m[0] for m in members) / len(members)
            mean_y = sum(m[1] for m in members) / len(members)
            new_centroids.append((mean_x, mean_y))

        # converged: no centroid moved, stop early
        if new_centroids == centroids:
            centroids = new_centroids
            break
        centroids = new_centroids

    return centroids, assignments

points = [(1, 1), (1.5, 2), (3, 4), (5, 7), (3.5, 5), (4.5, 5), (3.5, 4.5)]
centroids, assignments = kmeans(points, k=2, initial_centroids=[(1, 1), (1.5, 2)])
print(centroids)
print(assignments)

Complexity

O(n · k · i · d) time, where n is the number of points, k the number of clusters, i the number of iterations until convergence, and d the number of dimensions — every point is compared against every centroid on every iteration. O(n + k) space for the assignments and centroids.

When to Use It / Tradeoffs

• Needs k specified upfront — no way to ask "how many clusters" for you.
• Sensitive to initial centroid placement — a bad random start can converge to a visibly worse local optimum (see the Pitfalls section for the standard fix).
• Assumes roughly spherical, similar-sized clusters — struggles badly on elongated or nested shapes.
• No native concept of noise — every point is forced into some cluster, even an obvious outlier.
• Not deterministic in general (random initialization) — though the run above is, since the initial centroids were fixed rather than randomized.
• Fast and simple, which is why it's usually the first thing tried on numeric, roughly round clusters at moderate scale.

DBSCAN

DBSCAN (Density-Based Spatial Clustering of Applications with Noise) classifies every point as core (at least min_samples points, including itself, within distance eps), border (not core itself, but within eps of a core point), or noise (neither). Clusters are grown by chaining density-reachable core points together: start at an unvisited core point, claim every point in its eps-neighborhood into the same cluster, and for any of those that are themselves core, recursively pull in their neighborhoods too. Points that end up reachable only as border points get claimed but don't expand the cluster further.

Unlike K-Means, DBSCAN never needs a target number of clusters — the density structure of the data determines how many clusters exist — and it explicitly labels outliers as noise instead of forcing them into a cluster.

Worked Example

Nine 2D points: a tight group of four near (1,1), one extra point A5(2.3,1.2) sitting just off that group, a tight group of three near (5,5), and one far-away outlier at (10,10). eps=1.0, min_samples=3.

A5 is the interesting one: it only has 2 neighbors within eps (itself and A4), so it fails the core-point test on its own — but since it falls inside core point A4's neighborhood, it gets claimed as a border point of cluster 0 without ever expanding the cluster further. N1 has no neighbors at all within eps and is never claimed by any core point's neighborhood, so it stays noise.

import math

def euclidean_distance(a, b):
    return math.sqrt(sum((a[i] - b[i]) ** 2 for i in range(len(a))))

def region_query(points, idx, eps):
    return [j for j in range(len(points)) if euclidean_distance(points[idx], points[j]) <= eps]

def dbscan(points, eps, min_samples):
    n = len(points)
    # None = not yet visited, -1 = noise, 0..k = cluster id
    labels = [None] * n
    neighbors = [region_query(points, i, eps) for i in range(n)]
    # a point is a core point if enough points (including itself) fall within eps
    is_core = [len(neighbors[i]) >= min_samples for i in range(n)]

    next_cluster_id = -1
    for i in range(n):
        if labels[i] is not None:
            continue
        if not is_core[i]:
            # provisionally noise -- may still be claimed as a border point later
            labels[i] = -1
            continue

        next_cluster_id += 1
        labels[i] = next_cluster_id
        # grow the cluster outward from this core point
        seeds = list(neighbors[i])
        pos = 0
        while pos < len(seeds):
            j = seeds[pos]
            pos += 1
            if labels[j] == -1:
                # was noise -- claim it as a border point of this cluster
                labels[j] = next_cluster_id
            if labels[j] is not None:
                continue
            labels[j] = next_cluster_id
            if is_core[j]:
                # j is also core -- its neighborhood extends the cluster further
                for k in neighbors[j]:
                    if k not in seeds:
                        seeds.append(k)

    return labels, is_core

points = [
    (1.0, 1.0), (1.2, 1.5), (0.8, 1.3), (1.5, 1.0), (2.3, 1.2),
    (5.0, 5.0), (5.2, 5.3), (4.8, 4.7),
    (10.0, 10.0),
]
labels, is_core = dbscan(points, eps=1.0, min_samples=3)
print(labels)
print(is_core)

Complexity

O(n²) time as implemented above — region_query does a linear scan over all n points for each of the n points, since it's a plain distance check with no spatial index. With a spatial index (a k-d tree or grid, not implemented here since that would pull in extra machinery), neighbor queries drop to roughly O(log n) each, for O(n log n) overall. O(n) space for labels and cached neighbor lists.

When to Use It / Tradeoffs

• Doesn't need k specified upfront — the number of clusters falls out of the density structure.
• Explicitly labels outliers as noise instead of forcing them into a cluster.
• Handles non-convex, oddly-shaped clusters that K-Means can't (density-reachability chains through curves and irregular shapes fine).
• Deterministic given a fixed eps/min_samples and a fixed processing order (unlike K-Means's random-init non-determinism).
• eps and min_samples are hard to pick well and results are sensitive to them (see Pitfalls).
• Struggles when clusters have very different densities — a single global eps can't be simultaneously right for a dense cluster and a sparse one.

Hierarchical (Agglomerative) Clustering

Start with every point as its own singleton cluster. Repeatedly find the two closest clusters under some linkage rule and merge them into one, until only the target number of clusters remains (or, taken all the way, until every point has merged into one giant cluster). The full merge sequence from n singletons down to 1 cluster is a dendrogram — a tree of merges, ordered by the distance at which each merge happened — and you can "cut" it at any number of clusters after the fact, not just the one you originally asked for.

The linkage rule defines "distance between two clusters" (which isn't a single point-to-point distance once a cluster has more than one member): single linkage uses the distance between the two clusters' closest points, complete linkage the distance between their farthest points, average linkage the mean of every cross-cluster pairwise distance. The worked example and code below use single linkage, the simplest of the three to compute and trace by hand.

Worked Example

Six 2D points forming two visually obvious pairs-of-pairs plus a lone pair: P1(1,1), P2(1.5,1.5), P3(5,5), P4(5.5,5), P5(1,4), P6(1.2,4.5). Merging down to a single cluster, using single linkage:

The three tightest pairs merge first (steps 1-3, all under distance 0.71), exactly matching the visual grouping. Cutting the dendrogram at "2 clusters remaining" (right after step 4) recovers {P3,P4} and {P1,P2,P5,P6} — note that P1,P2 (near the origin) and P5,P6 (upper-left) end up in the same 2-cluster split, because with only 2 clusters allowed, single linkage's "closest pair" rule merges whichever two clusters are nearest at each step regardless of the original three visual groupings, and by step 4 {P3,P4} is still the odd one out at distance 3.833 from everything else.

import math

def euclidean_distance(a, b):
    return math.sqrt(sum((a[i] - b[i]) ** 2 for i in range(len(a))))

def single_linkage_distance(points, cluster_a, cluster_b):
    # single linkage: distance between two clusters is the distance
    # between their single closest pair of points
    return min(
        euclidean_distance(points[i], points[j])
        for i in cluster_a
        for j in cluster_b
    )

def agglomerative_clustering(points, num_clusters):
    # every point starts as its own cluster
    clusters = [[i] for i in range(len(points))]
    merge_log = []

    while len(clusters) > num_clusters:
        # find the two closest clusters by linkage distance
        best_pair = None
        best_distance = None
        for a in range(len(clusters)):
            for b in range(a + 1, len(clusters)):
                d = single_linkage_distance(points, clusters[a], clusters[b])
                if best_distance is None or d < best_distance:
                    best_distance = d
                    best_pair = (a, b)

        a, b = best_pair
        merged = clusters[a] + clusters[b]
        merge_log.append((list(clusters[a]), list(clusters[b]), best_distance))

        # replace the two merged clusters with the single merged one
        remaining = [c for idx, c in enumerate(clusters) if idx not in (a, b)]
        remaining.append(merged)
        clusters = remaining

    return clusters, merge_log

points = [(1, 1), (1.5, 1.5), (5, 5), (5.5, 5), (1, 4), (1.2, 4.5)]
clusters, merge_log = agglomerative_clustering(points, num_clusters=1)
print(clusters)
for a, b, d in merge_log:
    print(a, b, round(d, 3))

Complexity

O(n³) time as implemented above: n-1 merges happen in total, and each merge step recomputes the linkage distance between every remaining pair of clusters from scratch — an O(n²) scan per merge, times O(n) merges. This is the naive version; production implementations avoid recomputing every pairwise distance after each merge by maintaining a priority queue of candidate merges and updating only the entries touched by the merge, which gets the whole process down to O(n² log n). O(n²) space for the full pairwise distance matrix.

When to Use It / Tradeoffs

• Doesn't need k specified upfront — pick any cut of the dendrogram after the fact, or let the merge distances themselves suggest a natural cutoff.
• Deterministic — no random initialization step, unlike K-Means.
• Single linkage in particular can follow non-convex, chain-like shapes; complete/average linkage behave more like K-Means, biased toward compact clusters.
• No explicit noise handling — like K-Means, every point ends up in some cluster.
• O(n²) or worse complexity (even with the smarter bookkeeping) makes it impractical much past a few thousand points — see Pitfalls.

Union-Find: Clustering by Explicit Links

Union-Find is a fundamentally different approach to "clustering" than the three above: it never computes a distance or similarity between anything. Instead, you're given (or you discover via some separate, upstream check) explicit pairwise facts — "these two items belong together" — and Union-Find's job is purely to chain those pairwise facts into connected groups, transitively, in near-O(1) amortized time per operation. Every item starts in its own singleton group; union(a, b) merges two items' groups together; find(x) answers "which group is x in" by walking to that group's root.

The concrete example in this problem set is #9, Cluster Near-Duplicate Queries: an upstream embedding-similarity check has already flagged certain pairs of incoming support queries as similar enough to be near-duplicates, and the problem is to group all queries into clusters of mutually-connected duplicates (transitively — if q1~q2 and q2~q3 are both flagged, q1, q2, and q3 all land in one cluster even though q1~q3 was never directly flagged) so only one representative per cluster needs a fresh answer. See that page for the full worked trace over 8 nodes; it isn't repeated here.

Two optimizations make the near-O(1) bound possible: path compression (every node visited while walking up to the root during a find gets re-pointed straight at the root) and union by rank (always attach the shorter tree under the taller one, so trees never get needlessly deep).

class UnionFind:
    def __init__(self, n):
        self.parent = list(range(n))
        self.rank = [0] * n

    def find(self, x):
        # path compression: point every visited node straight at the root
        if self.parent[x] != x:
            self.parent[x] = self.find(self.parent[x])
        return self.parent[x]

    def union(self, a, b):
        root_a, root_b = self.find(a), self.find(b)
        if root_a == root_b:
            return False
        # union by rank: attach the shorter tree under the taller one
        if self.rank[root_a] < self.rank[root_b]:
            root_a, root_b = root_b, root_a
        self.parent[root_b] = root_a
        if self.rank[root_a] == self.rank[root_b]:
            self.rank[root_a] += 1
        return True

Complexity: amortized O(α(n)) per find/union call with both optimizations — effectively O(1); O(n) space for the parent and rank arrays.

Comparison

Pitfalls