← back

Implement Hierarchical Clustering (Agglomerative)

#364 · Machine Learning · Medium

⊣ Solve on deep-ml.com

Problem

Implement agglomerative (bottom-up) hierarchical clustering. Start with each data point as its own cluster, then iteratively merge the two closest clusters until the desired number of clusters is reached. Support single, complete, and average linkage.

Solution

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
import numpy as np

def agglomerative_clustering(
    X: np.ndarray, n_clusters: int, linkage: str = "single"
) -> np.ndarray:
    n = X.shape[0]
    # Initialize each point as its own cluster
    clusters = {i: [i] for i in range(n)}
    labels = np.arange(n)

    # Precompute pairwise distances
    dists = np.sum((X[:, None] - X[None, :]) ** 2, axis=2)
    np.fill_diagonal(dists, np.inf)

    current_clusters = set(range(n))

    while len(current_clusters) > n_clusters:
        # Find the two closest clusters
        min_dist = np.inf
        merge_a, merge_b = -1, -1
        cluster_list = sorted(current_clusters)
        for i in range(len(cluster_list)):
            for j in range(i + 1, len(cluster_list)):
                ci, cj = cluster_list[i], cluster_list[j]
                pts_i, pts_j = clusters[ci], clusters[cj]
                if linkage == "single":
                    d = min(dists[a, b] for a in pts_i for b in pts_j)
                elif linkage == "complete":
                    d = max(dists[a, b] for a in pts_i for b in pts_j)
                else:  # average
                    d = np.mean([dists[a, b] for a in pts_i for b in pts_j])
                if d < min_dist:
                    min_dist = d
                    merge_a, merge_b = ci, cj

        # Merge clusters
        clusters[merge_a] = clusters[merge_a] + clusters[merge_b]
        del clusters[merge_b]
        current_clusters.remove(merge_b)

    # Assign labels
    result = np.zeros(n, dtype=int)
    for label, (_, members) in enumerate(sorted(clusters.items())):
        for m in members:
            result[m] = label
    return result

Explanation

  1. Start with n singleton clusters, one per data point.
  2. At each step, compute inter-cluster distances using the chosen linkage: single (minimum pairwise), complete (maximum pairwise), or average (mean pairwise).
  3. Merge the two closest clusters into one.
  4. Repeat until only n_clusters clusters remain and assign labels.

Complexity

  • Time: O(n^3) since we check all cluster pairs at each of the O(n) merging steps
  • Space: O(n^2) for the distance matrix
← #363#365 →