k-Means Clustering & Lloyd's Optimization | AI & Machine Learning Engineering

Clustering partitions unlabelled data points into discrete subsets such that instances within the same cluster exhibit high mutual similarity, while instances across distinct clusters remain geometrically separated.

The canonical centroid-based clustering algorithm, k-Means, partitions continuous Euclidean feature space into a Voronoi diagram by minimizing the Within-Cluster Sum of Squares (WCSS).

Computational ComplexityEfficiency

O(N * K * d * I)

Linear scalability in samples N, clusters K, and dimensions d across I iterations.

1. Mathematical Formulation & The Inertia Objective#

Let $X = \{x_1, x_2, \dots, x_N\} \subset \mathbb{R}^d$ be an unlabelled dataset. We seek to partition $X$ into $K$ disjoint clusters $C = \{C_1, C_2, \dots, C_K\}$ with cluster prototypes (centroids) $\mu = \{\mu_1, \mu_2, \dots, \mu_K\} \subset \mathbb{R}^d$ .

Architecture Diagram

Rendering diagram...

We represent cluster membership via binary assignment indicator variables $r_{i,k} \in \{0, 1\}$ :

r_{i,k} = \begin{cases} 1 & \text{if } x_i \in C_k \\ 0 & \text{otherwise} \end{cases}, \quad \text{subject to} \quad \sum_{k=1}^K r_{i,k} = 1, \quad \forall i

The Inertia objective function measures total within-cluster variance:

J(r, \mu) = \sum_{i=1}^N \sum_{k=1}^K r_{i,k} \|x_i - \mu_k\|_2^2

2. Derivation of Lloyd's Coordinate Descent Algorithm#

Optimizing $J(r, \mu)$ simultaneously over discrete assignments $r$ and continuous centroids $\mu$ is an NP-hard combinatorial problem. Lloyd's Algorithm solves this using alternating coordinate descent.

Step 1: Optimize Assignments $r$ with Centroids $\mu$ Fixed

Hold centroids $\mu_k$ constant and minimize $J$ with respect to binary indicators $r_{i,k}$ :

\min_{r} \sum_{i=1}^N \sum_{k=1}^K r_{i,k} \|x_i - \mu_k\|_2^2 \quad \text{subject to} \quad \sum_{k=1}^K r_{i,k} = 1

Because the terms for each data point $i$ are decoupled, we minimize each sample independently by assigning $x_i$ to its nearest centroid:

r_{i,k}^* = \begin{cases} 1 & \text{if } k = \arg\min_{j \in \{1, \dots, K\}} \|x_i - \mu_j\|_2^2 \\ 0 & \text{otherwise} \end{cases}

Step 2: Optimize Centroids $\mu$ with Assignments $r$ Fixed

Now hold assignments $r_{i,k}$ fixed and minimize $J$ with respect to continuous vectors $\mu_k$ .

Differentiate $J$ with respect to vector $\mu_k$ :

\frac{\partial J}{\partial \mu_k} = \frac{\partial}{\partial \mu_k} \sum_{i=1}^N r_{i,k} (x_i - \mu_k)^T (x_i - \mu_k)

Apply the vector calculus identity $\frac{\partial}{\partial u} (v - u)^T (v - u) = -2(v - u)$ :

\frac{\partial J}{\partial \mu_k} = -2 \sum_{i=1}^N r_{i,k} (x_i - \mu_k) = 0

Step 3: Solve for Optimal Centroid $\mu_k^*$

Distribute the summation:

\sum_{i=1}^N r_{i,k} x_i - \sum_{i=1}^N r_{i,k} \mu_k = 0 \implies \left( \sum_{i=1}^N r_{i,k} \right) \mu_k = \sum_{i=1}^N r_{i,k} x_i

Let $N_k = \sum_{i=1}^N r_{i,k} = |C_k|$ denote the number of points assigned to cluster $k$ :

\mu_k^* = \frac{1}{|C_k|} \sum_{i \in C_k} x_i

Proof of Guaranteed Convergence#

Because both the Assignment Step and Update Step monotonically decrease (or leave unchanged) the non-negative scalar objective:

J(r^{(t+1)}, \mu^{(t+1)}) \le J(r^{(t+1)}, \mu^{(t)}) \le J(r^{(t)}, \mu^{(t)})

And because the number of distinct partitions of $N$ points into $K$ clusters is finite ( $S(N, K) \le K^N$ ), the sequence must converge to a local minimum or saddle point in a finite number of iterations.

3. k-Means++: Probabilistic Seeding Derivation#

Standard random centroid initialization frequently gets trapped in poor local minima with arbitrarily bad competitive ratios.

Arthur and Vassilvitskii (2007) introduced k-Means++, which initializes centroids sequentially by spreading them apart probabilistically.

Step 1: Initial Random Seed

Choose the first centroid $\mu_1$ uniformly at random from dataset $X$ .

Step 2: Distance to Nearest Existing Centroid

For each data point $x_i \in X$ , compute its shortest distance to any already chosen centroid:

D(x_i) = \min_{j \in \{1, \dots, m\}} \|x_i - \mu_j\|_2

Step 3: Probabilistic Sampling

Sample the next centroid $\mu_{m+1} = x$ from $X$ with probability proportional to the squared distance:

P(x) = \frac{D(x)^2}{\sum_{i=1}^N D(x_i)^2}

Repeat Steps 2 and 3 until exactly $K$ centroids are chosen.

Theoretical Guarantee

k-Means++ guarantees that the expected inertia

\mathbb{E}[J]

of the resulting clustering satisfies an

O(\log K)

approximation bound relative to the global optimal clustering

J^*

\mathbb{E}[J] \le 8(\ln K + 2) J^*

4. From-Scratch Implementation: Zero Dependencies#

The following implementations construct the complete k-Means Clustering with k-Means++ Initialization from scratch in Python and Rust.

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
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
import math
import random
from typing import List, Tuple

class KMeansFromScratch:
    def __init__(self, k: int = 3, max_iter: int = 300, tol: float = 1e-4):
        self.k = k
        self.max_iter = max_iter
        self.tol = tol
        self.centroids: List[List[float]] = []

    @staticmethod
    def _euclidean_dist_sq(a: List[float], b: List[float]) -> float:
        return sum((x_i - y_i) ** 2 for x_i, y_i in zip(a, b))

    def _kmeans_plus_plus_init(self, X: List[List[float]]) -> List[List[float]]:
        num_samples = len(X)
        # 1. First centroid selected uniformly at random
        centroids = [list(random.choice(X))]

        # 2. Select remaining K-1 centroids
        for _ in range(1, self.k):
            # Compute D(x)^2 for all samples
            d_sq = []
            for sample in X:
                min_dist = min(self._euclidean_dist_sq(sample, c) for c in centroids)
                d_sq.append(min_dist)

            total_dist = sum(d_sq)
            # Sample with probability proportional to D(x)^2
            r = random.uniform(0.0, total_dist)
            cumulative = 0.0
            next_centroid = X[-1]
            for i, dist in enumerate(d_sq):
                cumulative += dist
                if cumulative >= r:
                    next_centroid = X[i]
                    break
            centroids.append(list(next_centroid))

        return centroids

    def fit(self, X: List[List[float]]) -> None:
        random.seed(42)
        num_samples = len(X)
        num_features = len(X[0])
        self.centroids = self._kmeans_plus_plus_init(X)

        for _ in range(self.max_iter):
            # Assignment Step: partition points into K clusters
            clusters: List[List[List[float]]] = [[] for _ in range(self.k)]
            for sample in X:
                nearest_idx = min(
                    range(self.k),
                    key=lambda idx: self._euclidean_dist_sq(sample, self.centroids[idx])
                )
                clusters[nearest_idx].append(sample)

            # Update Step: recompute cluster centroids as arithmetic mean
            new_centroids: List[List[float]] = []
            max_shift = 0.0

            for k in range(self.k):
                if not clusters[k]:
                    # Handle empty cluster by keeping old centroid
                    new_centroids.append(self.centroids[k])
                    continue

                cluster_size = len(clusters[k])
                centroid = [
                    sum(sample[j] for sample in clusters[k]) / cluster_size
                    for j in range(num_features)
                ]
                shift = math.sqrt(self._euclidean_dist_sq(centroid, self.centroids[k]))
                max_shift = max(max_shift, shift)
                new_centroids.append(centroid)

            self.centroids = new_centroids
            if max_shift < self.tol:
                break

    def predict(self, X: List[List[float]]) -> List[int]:
        predictions = []
        for sample in X:
            nearest_idx = min(
                range(self.k),
                key=lambda idx: self._euclidean_dist_sq(sample, self.centroids[idx])
            )
            predictions.append(nearest_idx)
        return predictions

if __name__ == "__main__":
    X = [
        [1.0, 1.0], [1.5, 2.0], [1.2, 0.8],  # Cluster 0
        [8.0, 8.0], [8.5, 9.0], [9.0, 8.2],  # Cluster 1
        [1.0, 8.0], [1.2, 8.5], [0.8, 9.0],  # Cluster 2
    ]

    kmeans = KMeansFromScratch(k=3)
    kmeans.fit(X)
    preds = kmeans.predict(X)
    print("Final Centroids:", kmeans.centroids)
    print("Cluster Assignments:", preds)

5. Active Knowledge Checks#

Concept check

In Lloyd's algorithm for k-Means, why does the centroid update step set mu_k equal to the exact arithmetic mean of assigned points?

Concept check

What is the foundational theoretical advantage of k-Means++ probabilistic initialization over uniform random centroid initialization?