Implement k-Means from Scratch

Math Foundations

Given unlabeled dataset $X \in \mathbb{R}^{N \times D}$, k-Means Clustering partitions data into $K$ distinct clusters $\mathcal{C} = {C_1, C_2, \dots, C_K}$ with centroids $\mu = {\mu_1, \mu_2, \dots, \mu_K}$.

The goal is to minimize Within-Cluster Sum of Squares (Inertia):

$$\mathcal{J} = \sum_{k=1}^K \sum_{x \in C_k} | x - \mu_k |^2$$

The algorithm alternates between two steps until centroid positions stop changing:

1. Assignment Step: Assign each sample $x_i$ to the closest centroid:

$$c_i = \arg\min_{k} | x_i - \mu_k |^2$$

2. Update Step: Recompute centroid $\mu_k$ as the mean of assigned points:

$$\mu_k = \frac{1}{|C_k|} \sum_{x_i \in C_k} x_i$$

NumPy Implementation from Scratch

import numpy as np


class KMeansFromScratch:

  def __init__(self, k=3, max_iters=100, tol=1e-4):
    self.k = k
    self.max_iters = max_iters
    self.tol = tol
    self.centroids = None

  def fit(self, X):
    n_samples, n_features = X.shape

    # 1. Randomly initialize K centroids from data points
    random_indices = np.random.choice(n_samples, self.k, replace=False)
    self.centroids = X[random_indices]

    for _ in range(self.max_iters):
      # 2. Assignment Step: Compute distances to all centroids
      # Distances shape: [n_samples, k]
      distances = np.linalg.norm(
          X[:, np.newaxis, :] - self.centroids[np.newaxis, :, :], axis=2
      )
      labels = np.argmin(distances, axis=1)

      # 3. Update Step: Recompute centroids as cluster means
      new_centroids = np.zeros((self.k, n_features))
      for k_idx in range(self.k):
        cluster_points = X[labels == k_idx]
        if len(cluster_points) > 0:
          new_centroids[k_idx] = np.mean(cluster_points, axis=0)
        else:
          # Handle empty cluster: re-initialize to a random data point
          new_centroids[k_idx] = X[np.random.choice(n_samples)]

      # 4. Check Convergence
      centroid_shift = np.sum(np.abs(new_centroids - self.centroids))
      self.centroids = new_centroids
      if centroid_shift < self.tol:
        break

    return labels

  def predict(self, X):
    distances = np.linalg.norm(
        X[:, np.newaxis, :] - self.centroids[np.newaxis, :, :], axis=2
    )
    return np.argmin(distances, axis=1)

Matrix Distance Optimization

Notice the distance calculation using NumPy broadcasting:

X[:, np.newaxis, :] - self.centroids[np.newaxis, :, :]

• X shape: [N, 1, D]
• centroids shape: [1, K, D]
• Difference shape: [N, K, D]

This computes all $N \times K$ pairwise Euclidean distances in parallel without explicit Python loops!

Say this out loud

Implementing k-Means from scratch alternates between assigning data points to the nearest centroid and recomputing centroids as cluster means. Vectorized distance broadcasting computes pairwise Euclidean distances efficiently across all points and centroids. Random initialization can be improved using k-Means++ to spread out initial centroids.

Followups to expect

1. What is k-Means++ Initialization? Choosing the first centroid at random, then picking subsequent centroids with probability proportional to their squared distance from nearest existing centroids, reducing convergence time and local minima risks.
2. How do you choose optimal K using the Elbow Method? Plotting Inertia against candidate values of K and selecting the point where the rate of inertia decrease sharpens into an elbow shape.