K-Means Clustering From Scratch in Python: The Algorithm, Math, and Code Explained

K-Means Clustering is an unsupervised machine learning algorithm designed to partition an unlabeled dataset of $n$ observations into $K$ distinct, non-overlapping groups (clusters). Because real-world data rarely comes with pre-assigned target labels, K-Means discovers the natural geometric groupings latent within the feature space purely by measuring physical proximity.

1. Key K-Means Concepts & Notation

Before stepping through the optimization formulas, let us establish the formal mathematical notation and terminology used throughout clustering theory:

2. Mathematical Derivations & Proof of Convergence

K-Means operates as an optimization algorithm that minimizes a single objective function: the Within-Cluster Sum of Squares (WCSS), also known as Inertia:

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

Where $x_i$ is a $m$-dimensional feature vector, $C_k$ is the set of points assigned to cluster $k$, and $\mu_k$ is the centroid of cluster $k$. K-Means minimizes this objective using an alternating two-step optimization known as Lloyd's Algorithm (a hard-assignment instance of Expectation-Maximization):

Step 1: The Assignment Phase (Expectation / E-Step)

Holding the centroids $\mu_1, \dots, \mu_K$ fixed, each observation $x_i$ is assigned to its nearest centroid according to squared Euclidean distance:

$$c_i = \arg\min_{k \in \{1, \dots, K\}} \| x_i - \mu_k \|^2$$

Because each point independently chooses the centroid that minimizes its individual contribution $\|x_i - \mu_k\|^2$, this step strictly decreases (or holds constant) the total objective $J$.

Step 2: The Centroid Update Phase (Maximization / M-Step)

Holding the cluster assignments $C_1, \dots, C_K$ fixed, we recompute each centroid $\mu_k$ to minimize $J$. To prove why the arithmetic mean is the optimal choice, we compute the partial derivative of $J$ with respect to centroid $\mu_k$ and set it to zero:

$$\frac{\partial J}{\partial \mu_k} = \frac{\partial}{\partial \mu_k} \sum_{x_i \in C_k} \| x_i - \mu_k \|^2 = -2 \sum_{x_i \in C_k} (x_i - \mu_k) = 0$$

$$\sum_{x_i \in C_k} x_i - \sum_{x_i \in C_k} \mu_k = 0 \implies \sum_{x_i \in C_k} x_i - |C_k| \mu_k = 0$$

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

This proves analytically that updating $\mu_k$ to the component-wise arithmetic mean of its assigned observations achieves the unique, global minimum of the squared error for that cluster partition.

3. Worked Numerical Trace (By Hand with 12 Points)

To see the mathematics operating in practice, let us trace an exact execution on a 2D dataset of 12 points grouped into 3 distinct spatial regions:

• Bottom-Left Group (4 points): $p_1=(1.0, 2.0), p_2=(1.5, 1.8), p_3=(1.0, 0.6), p_4=(2.0, 1.0)$
• Top-Right Group (4 points): $p_5=(9.0, 11.0), p_6=(8.0, 8.0), p_7=(9.0, 10.0), p_8=(10.0, 9.0)$
• Bottom-Right Group (4 points): $p_9=(8.0, 2.0), p_{10}=(9.0, 3.0), p_{11}=(10.0, 2.0), p_{12}=(9.0, 1.0)$

Suppose we initialize the three centroids directly on observations $p_1$, $p_5$, and $p_9$:

$$\mu_1 = (1.0, 2.0), \quad \mu_2 = (9.0, 11.0), \quad \mu_3 = (8.0, 2.0)$$

Iteration 1: Assignment & Centroid Update

Each of the 12 points computes its Euclidean distance to $\mu_1$, $\mu_2$, and $\mu_3$. Because the three groups are widely separated, points $p_1..p_4$ assign to $\mu_1$, points $p_5..p_8$ assign to $\mu_2$, and points $p_9..p_{12}$ assign to $\mu_3$. We then recompute the centroids:

• $$\mu_1 = \left( \frac{1.0 + 1.5 + 1.0 + 2.0}{4}, \frac{2.0 + 1.8 + 0.6 + 1.0}{4} \right) = \left( \frac{5.5}{4}, \frac{5.4}{4} \right) = (1.375, 1.350)$$
• $$\mu_2 = \left( \frac{9.0 + 8.0 + 9.0 + 10.0}{4}, \frac{11.0 + 8.0 + 10.0 + 9.0}{4} \right) = \left( \frac{36.0}{4}, \frac{38.0}{4} \right) = (9.000, 9.500)$$
• $$\mu_3 = \left( \frac{8.0 + 9.0 + 10.0 + 9.0}{4}, \frac{2.0 + 3.0 + 2.0 + 1.0}{4} \right) = \left( \frac{36.0}{4}, \frac{8.0}{4} \right) = (9.000, 2.000)$$

Iteration 2: Convergence Check & Inertia Calculation

In Iteration 2, every point recalculates distances against the updated centroids. Zero points switch clusters. Because no assignments change, the centroids remain identical, and the algorithm terminates.

Now calculate the exact Within-Cluster Sum of Squares ($J$):

• Cluster 1 WCSS: $(1.0-1.375)^2 + (2.0-1.35)^2 + (1.5-1.375)^2 + (1.8-1.35)^2 + (1.0-1.375)^2 + (0.6-1.35)^2 + (2.0-1.375)^2 + (1.0-1.35)^2 = 1.6975$
• Cluster 2 WCSS: $(9-9)^2 + (11-9.5)^2 + (8-9)^2 + (8-9.5)^2 + (9-9)^2 + (10-9.5)^2 + (10-9)^2 + (9-9.5)^2 = 5.0000$
• Cluster 3 WCSS: $(8-9)^2 + (2-2)^2 + (9-9)^2 + (3-2)^2 + (10-9)^2 + (2-2)^2 + (9-9)^2 + (1-2)^2 = 6.3000$

$$\text{Total Inertia } J = 1.6975 + 5.0000 + 6.3000 = 12.9975$$

4. From-Scratch Python Implementation (Vectorized NumPy)

Below is a fully vectorized, object-oriented implementation of K-Means built with NumPy. It incorporates tolerance checking, inertia calculation, and safeguards against empty cluster divisions:

import numpy as np

class KMeansScratch:
    """
    K-Means clustering implemented from scratch using NumPy.
    Minimizes Within-Cluster Sum of Squares (WCSS / Inertia).
    """
    def __init__(self, k=3, max_iters=100, tol=1e-4, random_state=42):
        self.k = k
        self.max_iters = max_iters
        self.tol = tol
        self.random_state = random_state
        self.centroids = None
        self.labels = None
        self.inertia_ = None

    def _init_centroids(self, X, rng):
        # Pick k random distinct observations as starting centroids
        indices = rng.choice(len(X), size=self.k, replace=False)
        return X[indices].astype(float).copy()

    def fit(self, X):
        X = np.asarray(X, dtype=float)
        n_samples, n_features = X.shape
        rng = np.random.default_rng(self.random_state)

        self.centroids = self._init_centroids(X, rng)

        for iteration in range(self.max_iters):
            # 1. Assignment Step (E-step): compute distance matrix (n_samples, k)
            # Using Euclidean distance: ||x_i - mu_k||
            distances = np.linalg.norm(X[:, np.newaxis, :] - self.centroids[np.newaxis, :, :], axis=2)
            new_labels = np.argmin(distances, axis=1)

            # 2. Update Step (M-step): recalculate centroids as mean of assigned points
            new_centroids = np.zeros_like(self.centroids)
            for k_idx in range(self.k):
                cluster_members = X[new_labels == k_idx]
                if len(cluster_members) > 0:
                    new_centroids[k_idx] = cluster_members.mean(axis=0)
                else:
                    # Safeguard against empty clusters: re-seed to the point farthest from any centroid
                    furthest_idx = np.argmax(np.min(distances, axis=1))
                    new_centroids[k_idx] = X[furthest_idx]

            # 3. Check for convergence (centroid movement < tolerance)
            centroid_shift = np.linalg.norm(new_centroids - self.centroids)
            self.centroids = new_centroids
            self.labels = new_labels

            if centroid_shift < self.tol:
                break

        # Compute final WCSS (Inertia)
        self.inertia_ = sum(
            np.sum((X[self.labels == k_idx] - self.centroids[k_idx]) ** 2)
            for k_idx in range(self.k)
        )
        return self

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


# Verification test on the 12-point toy dataset
if __name__ == "__main__":
    X = np.array([
        [1.0, 2.0], [1.5, 1.8], [1.0, 0.6], [2.0, 1.0],      # bottom-left group
        [9.0, 11.0], [8.0, 8.0], [9.0, 10.0], [10.0, 9.0],   # top-right group
        [8.0, 2.0], [9.0, 3.0], [10.0, 2.0], [9.0, 1.0],     # bottom-right group
    ])

    model = KMeansScratch(k=3, max_iters=100, random_state=42)
    # Seed centroids directly with initial indices to trace exact hand calculation
    model.centroids = np.array([[1.0, 2.0], [9.0, 11.0], [8.0, 2.0]])
    model.fit(X)

    print("Final Centroids:\n", np.round(model.centroids, 4))
    print("Cluster Labels:", model.labels)
    print(f"Final WCSS (Inertia): {model.inertia_:.4f}")

5. Production Scikit-Learn Pipeline & K-Means++

While naive K-Means picks starting centroids uniformly at random, a poor initial draw can trap the algorithm in an inferior local minimum. To solve this, Arthur and Vassilvitskii (2007) introduced K-Means++, which is the default in Scikit-Learn (init='k-means++').

How K-Means++ Works

1. Choose the first centroid $\mu_1$ uniformly at random from the dataset.
2. For every remaining point $x$, compute $D(x)$, the Euclidean distance to the closest centroid already selected.
3. Select the next centroid $\mu_{next}$ by sampling points with probability proportional to the squared distance:

$$P(x) = \frac{D(x)^2}{\sum_{x' \in X} D(x')^2}$$

Repeat steps 2 and 3 until $K$ centroids are chosen. By favoring points that are physically far away from existing centroids, K-Means++ spreads the starting seeds across the data manifold, guaranteeing an expected approximation ratio of $O(\log K)$ compared to the global optimal clustering.

import numpy as np
from sklearn.cluster import KMeans
from sklearn.preprocessing import StandardScaler
from sklearn.pipeline import Pipeline
from sklearn.metrics import silhouette_score, silhouette_samples

# 1. Real-world dataset simulation
X = np.array([
    [1.0, 2.0], [1.5, 1.8], [1.0, 0.6], [2.0, 1.0],
    [9.0, 11.0], [8.0, 8.0], [9.0, 10.0], [10.0, 9.0],
    [8.0, 2.0], [9.0, 3.0], [10.0, 2.0], [9.0, 1.0],
])

# 2. Production Pipeline with StandardScaler & K-Means++
# StandardScaler ensures features contribute equally to Euclidean distances
pipeline = Pipeline([
    ("scaler", StandardScaler()),
    ("kmeans", KMeans(n_clusters=3, init="k-means++", n_init=10, random_state=42))
])

pipeline.fit(X)
kmeans_model = pipeline.named_steps["kmeans"]
X_scaled = pipeline.named_steps["scaler"].transform(X)

# 3. Model Evaluation: Inertia & Silhouette Score
labels = kmeans_model.labels_
sil_score = silhouette_score(X_scaled, labels)

print("Cluster Labels:           ", labels)
print(f"Inertia (Scaled WCSS):    {kmeans_model.inertia_:.4f}")
print(f"Silhouette Coefficient:   {sil_score:.4f} (Close to +1 indicates excellent separation)")

# 4. Assessing multiple values of K (Elbow & Silhouette Analysis)
print("\n--- Evaluating K from 2 to 5 ---")
for k in range(2, 6):
    km = KMeans(n_clusters=k, init="k-means++", n_init=10, random_state=42)
    pred_labels = km.fit_predict(X_scaled)
    sil = silhouette_score(X_scaled, pred_labels)
    print(f"K={k} | WCSS (Inertia): {km.inertia_:7.3f} | Silhouette Score: {sil:.3f}")

6. How to Choose K: The Elbow Method vs. Silhouette Analysis

Because K-Means requires specifying $K$ upfront, practitioners use two complementary validation techniques to determine the optimal cluster count:

The Elbow Method

The Elbow Method plots WCSS (Inertia) against increasing values of $K$. As $K$ increases, inertia naturally decreases (reaching 0 when $K=n$, where every point is its own centroid). The ideal cluster count corresponds to the 'elbow' inflection point — where adding another cluster yields diminishing marginal returns in variance reduction.

The Silhouette Coefficient

When the elbow curve is smooth and ambiguous, the Silhouette Score provides an objective measure of cluster cohesion and separation for each point $i$:

$$s(i) = \frac{b(i) - a(i)}{\max(a(i), b(i))}$$

• $a(i)$: The mean intra-cluster distance from point $i$ to all other points in the same cluster (measures compactness).
• $b(i)$: The mean nearest-cluster distance from point $i$ to all points in the closest neighboring cluster (measures separation).
• Interpretation: A score near $+1$ indicates point $i$ is well inside its cluster and far from neighbors; $0$ indicates on the decision border; $-1$ indicates point $i$ was assigned to the wrong cluster.

7. When K-Means Fails: 4 Structural Limitations

• Spherical Cluster Assumption: K-Means inherently assumes clusters are isotropic and spherical because it assigns points purely based on Euclidean distance to a single center. If true clusters are elongated ellipses (anisotropic), K-Means will split them incorrectly perpendicular to the long axis (use Gaussian Mixture Models instead).
• Unequal Cluster Sizes & Densities: When one cluster contains 10,000 points and a neighboring cluster contains 50 points, K-Means often pulls the centroid of the small cluster into the dense cloud to minimize total squared error.
• Non-Convex Geometries: Data shaped like concentric rings, interlocking spirals, or crescent moons cannot be separated by K-Means' linear Voronoi boundaries (use DBSCAN or Spectral Clustering).
• Sensitivity to Feature Scales & Outliers: Because squared Euclidean distance $(\Delta x)^2$ penalizes large deviations quadratically, unscaled features dominate distance calculations, and extreme outliers drag centroids far away from true cluster centers.

8. Summary & Key Takeaways

• Unsupervised Clustering: K-Means discovers hidden patterns in unlabeled data by grouping observations around $K$ spatial centroids.
• Lloyd's Two-Step Loop: Alternates between assigning points to the closest centroid and updating centroids to the mean of their members.
• Minimizes WCSS: Operates as coordinate descent on Within-Cluster Sum of Squares, guaranteed to converge in a finite number of steps.
• Always Scale Features: Standardization (StandardScaler) is mandatory so high-magnitude features do not distort distance calculations.
• Use K-Means++: Spreads initial seeds using $D^2$ distance weighting, avoiding sub-optimal local minima traps.