K-Means Clustering From Scratch
Implement the k-means clustering algorithm from scratch and return its final inertia.
Given a set of $n$ data points in $d$-dimensional space and a target number of clusters $k$, write a function
```python def kmeans_inertia(points, k): ... ```
that partitions the points into $k$ clusters by iteratively assigning points to their nearest centroid and updating centroids, then returns the within-cluster sum of squared Euclidean distances (the inertia) as a single float.
Deterministic initialization (required so the answer is unique): - centroid[i] = points[i] for i in 0..k-1 (the first $k$ points).
Then run Lloyd iterations until assignments stop changing (max 100 iterations): - Assign: each point → nearest centroid (Euclidean distance). - Update: each centroid → mean of its assigned points. If a cluster is empty, leave its centroid unchanged.
Inertia $= \sum_{\text{points } p} \|p - \text{assigned\_centroid}(p)\|^2$.
Constraints: - $1 \leq k \leq n \leq 10^4$ - $1 \leq d \leq 100$ - Use Euclidean distance.
Example
kmeans_inertia([[1,2],[1,4],[1,0],[4,2],[4,4],[4,0]], 2) -> 17.5
Centroids are initialized to the first 2 points [1,2] and [1,4]. Lloyd iterations then converge to a local minimum whose within-cluster sum of squared distances is $17.5$. Because initialization is fixed to the first $k$ points rather than optimized, the result is a deterministic local optimum, not necessarily the globally best clustering.
Hints
- The algorithm has just two steps that repeat: assign each point to the closest centroid, then recompute each centroid as the mean of its cluster. Think about why the objective must decrease at each step.
- For initialization, random selection can lead to poor local minima. K-means++ picks initial centroids proportional to squared distance from existing centroids, spreading them out.
- Convergence is guaranteed because each step decreases the objective $\sum \|x - \mu_i\|^2$ and there are finitely many possible assignments. Use a flag to detect when no assignments change.
Worked Solution
How to Think About It: K-means alternates two steps until convergence: (1) assign each point to its nearest centroid (Euclidean distance), and (2) update each centroid to the mean of its assigned points. Both steps decrease the objective $\sum_{i=1}^k \sum_{x \in C_i} \|x - \mu_i\|^2$ (the *inertia*), so Lloyd's algorithm converges to a local minimum in finitely many steps.
Contract for this problem: The judge calls kmeans_inertia(points, k) and expects a single float — the final inertia (within-cluster sum of squared Euclidean distances). Initialization is deterministic: centroid[i] = points[i] for i in 0..k-1 (the first k points). This is required so the returned number is unique. Do not use random or k-means++ init here, and do not return centroids/assignments — return the total.
Algorithm:
1. Set centroids to the first k points. 2. Repeat up to 100 iterations: - Assign: each point → nearest centroid. - Stop if assignments are unchanged from the previous iteration. - Update: each centroid → mean of its assigned points; if a cluster is empty, leave that centroid unchanged. 3. Inertia = sum over all points of $\|point - assigned\_centroid\|^2$.
Code:
```python import numpy as np
def kmeans_inertia(points, k): # Run k-means (Lloyd's algorithm) from scratch with deterministic # initialization (first k points as centroids) and return the final # within-cluster sum of squared Euclidean distances (inertia). pts = np.asarray(points, dtype=float) n, d = pts.shape
# Deterministic init: first k points centroids = pts[:k].copy()
assignments = None for _ in range(100): # Assign step: each point -> nearest centroid (Euclidean distance) dists = np.linalg.norm(pts[:, None, :] - centroids[None, :, :], axis=2) new_assignments = np.argmin(dists, axis=1)
if assignments is not None and np.array_equal(new_assignments, assignments): break assignments = new_assignments
# Update step: centroid -> mean of assigned points (empty -> unchanged) for i in range(k): mask = assignments == i if mask.any(): centroids[i] = pts[mask].mean(axis=0)
# Inertia: sum of squared distances of each point to its assigned centroid dists = np.linalg.norm(pts[:, None, :] - centroids[None, :, :], axis=2) final = np.argmin(dists, axis=1) diff = pts - centroids[final] inertia = float(np.sum(diff * diff)) return inertia ```
Complexity: - Time per iteration: $O(nkd)$ — all $n \times k$ pairwise distances in $d$ dimensions. - Space: $O(nd + kd)$. - Iterations bounded at 100 (typically far fewer before assignments stabilize).
Answer: With deterministic first-k-points initialization, run Lloyd iterations to convergence and return the within-cluster sum of squared distances as a float.
Intuition
K-means is a coordinate descent algorithm in disguise. The objective $\sum_{i=1}^k \sum_{x \in C_i} \|x - \mu_i\|^2$ depends on two sets of variables: the assignments and the centroids. The assign step optimizes over assignments with centroids fixed (each point goes to its nearest centroid). The update step optimizes over centroids with assignments fixed (the mean minimizes the sum of squared distances). Each step decreases the objective, and since there are finitely many possible partitions, the algorithm must terminate.
The practical lesson is that k-means is fast and simple, but it only finds local optima. In real applications -- customer segmentation, feature quantization, initializing Gaussian mixture models -- you always run it multiple times with different initializations and keep the best result. K-means++ largely solves the initialization problem, but understanding that k-means is fundamentally a local search algorithm is important for knowing when to trust its output and when to be skeptical.