K-means Clustering › Formula

Michael BrenndoerferAugust 2, 202520 min read
Machine Learning from ScratchAll chapters
Part 7 of 9Unsupervised Learning
Chapter 53K-means Clustering › Formula
Reading section 2 of 15 · Formula · 20 minSection 2 of 15 · 20 min

Book contents

Section 2 of 15 · 20 min read

K-means Clustering › FormulaOne section · Sections in this chapter2 / 15 ↓

K-means clustering is fundamentally about finding natural groupings in data by minimizing how spread out points are within each cluster. To understand the mathematics behind this intuitive goal, we'll build up from a simple question to the complete algorithm, revealing why each mathematical component is necessary and how they work together to solve the clustering problem.

The Core Question: What Makes a Good Clustering?

Imagine you have a collection of data points scattered across space, and you want to organize them into k groups. What would make one grouping better than another? Intuitively, we want points within the same group to be close together, tightly packed around some central location. This intuition leads us to the mathematical foundation of K-means.

For each cluster, we can represent its "center" by a single point called the centroid. Think of the centroid as the average position of all points in that cluster: it's the natural meeting point that minimizes the total distance to all cluster members. Once we have centroids, we can measure how "tight" each cluster is by calculating how far each point is from its centroid. The smaller these distances, the more compact and well-defined our clusters are.

This brings us to the fundamental objective: we want to find cluster assignments and centroid positions that minimize the total squared distance from each point to its assigned centroid. This is called the within-cluster sum of squares (WCSS) or inertia, and it's what K-means tries to minimize.

Building the Objective Function

Let's formalize this intuition step by step. Suppose we want to partition our dataset into k clusters, which we'll call C1,C2,…,CkC_1, C_2, \ldots, C_k. Each cluster CiC_i is simply a set of data points that belong together. For each cluster, we'll compute a centroid μi\mu_i that represents the cluster's center.

Now, how do we measure the quality of this clustering? We calculate the squared distance from each point to its cluster's centroid, then sum up all these squared distances across all clusters:

WCSS=∑i=1k∑x∈Ci∥x−μi∥2\text{WCSS} = \sum_{i=1}^{k} \sum_{x \in C_i} \|x - \mu_i\|^2

Let's unpack what each part means:

  • C1,C2,…,CkC_1, C_2, \ldots, C_k: The k clusters we're trying to find. Each cluster is a set of data points.
  • x∈Cix \in C_i: "For each data point x that belongs to cluster i." This notation tells us to consider every point in the cluster.
  • μi\mu_i: The centroid (mean position) of cluster i. This is the point that represents the cluster's center.
  • ∥x−μi∥2\|x - \mu_i\|^2: The squared Euclidean distance between point x and its cluster's centroid μi\mu_i. This measures how far the point is from its cluster center.
  • The double summation: The inner sum (∑x∈Ci\sum_{x \in C_i}) adds up distances for all points in cluster i. The outer sum (∑i=1k\sum_{i=1}^{k}) adds up the results across all k clusters.

Our goal is to find the cluster assignments and centroids that minimize this total squared distance:

min⁡C1,C2,…,Ck∑i=1k∑x∈Ci∥x−μi∥2\min_{C_1, C_2, \ldots, C_k} \sum_{i=1}^{k} \sum_{x \in C_i} \|x - \mu_i\|^2

Why squared distances? You might wonder why we square the distances rather than just using regular distances. There are two important reasons:

  1. Mathematical tractability: Squared distances make the optimization problem differentiable and easier to solve mathematically. The square function is smooth everywhere, which allows us to use calculus-based optimization techniques.

  2. Emphasis on outliers: Squaring distances penalizes points that are far from their centroids more heavily than points that are close. A point that's twice as far away contributes four times as much to the objective function. This encourages the algorithm to create tight, compact clusters.

Let's visualize what the WCSS objective function represents:

Out[2]:
Visualization
Scatter plot showing three clusters with centroids and dashed lines connecting points to their centroids to illustrate WCSS calculation.
Visual representation of the Within-Cluster Sum of Squares (WCSS). This diagram shows three clusters with their centroids (marked with red ''x'' symbols). The gray dashed lines connect each data point to its assigned centroid, and the length of these lines represents the distance we''re trying to minimize. The WCSS is the sum of all these squared distances across all clusters. Points closer to their centroids contribute less to WCSS, while points farther away contribute more (especially due to squaring). This visualization makes the abstract mathematical formula concrete: we''re literally summing up all the squared lengths of these gray lines.

Understanding Distance in Detail

Let's zoom in on the distance calculation itself. When we write ∥x−μi∥2\|x - \mu_i\|^2, we're computing the squared Euclidean distance between two points in d-dimensional space. If our data has d features, then each point x and each centroid μi\mu_i is a d-dimensional vector.

The squared Euclidean distance is calculated by summing the squared differences across all dimensions:

∥x−μi∥2=∑j=1d(xj−μi,j)2\|x - \mu_i\|^2 = \sum_{j=1}^{d} (x_j - \mu_{i,j})^2

The components in this expression have the following meanings:

  • dd: The number of features (dimensions) in our data. For example, if we're clustering customers by age and income, then d = 2.
  • xjx_j: The value of the j-th feature for data point x. For instance, x2x_2 might be the income of a customer.
  • μi,j\mu_{i,j}: The value of the j-th feature for centroid μi\mu_i. This is the average value of feature j across all points in cluster i.
  • (xj−μi,j)2(x_j - \mu_{i,j})^2: The squared difference between the point's feature value and the centroid's feature value for dimension j.

For example, if we have a 2D point x=(3,5)x = (3, 5) and a centroid μi=(1,2)\mu_i = (1, 2), the squared distance is:

∥x−μi∥2=(3−1)2+(5−2)2=4+9=13\|x - \mu_i\|^2 = (3-1)^2 + (5-2)^2 = 4 + 9 = 13

Let's visualize this distance calculation geometrically:

Out[3]:
Visualization
Geometric diagram showing Euclidean distance calculation as a right triangle with coordinates and Pythagorean theorem.
Geometric interpretation of Euclidean distance calculation. This diagram shows how the squared Euclidean distance between point x = (3, 5) and centroid μᵢ = (1, 2) is computed. The horizontal difference (3-1 = 2) and vertical difference (5-2 = 3) form the two legs of a right triangle. By the Pythagorean theorem, the distance is √(2² + 3²) = √13 ≈ 3.61. The squared distance is 13, which is what K-means actually uses in its objective function. This visualization shows that Euclidean distance is simply the Pythagorean theorem generalized to d dimensions: we square each coordinate difference, sum them up, and take the square root (though K-means skips the square root step).

This distance calculation treats all dimensions equally, which is why you should standardize the data before applying K-means. If one feature has values in the thousands while another has values between 0 and 1, the larger-scale feature will dominate the distance calculation and disproportionately influence the clustering.

Computing the Centroid: The Natural Center

Now that we understand how to measure distances, let's see how to compute the centroid of a cluster. The centroid is simply the mean (average) position of all points in the cluster. This choice isn't arbitrary: the mean is mathematically special because it's the point that minimizes the sum of squared distances to all other points in the set.

For cluster i, the centroid is calculated as:

μi=1∣Ci∣∑x∈Cix\mu_i = \frac{1}{|C_i|} \sum_{x \in C_i} x

Breaking this down:

  • ∣Ci∣|C_i|: The number of points in cluster i. The vertical bars denote the cardinality (size) of the set.
  • ∑x∈Cix\sum_{x \in C_i} x: The sum of all data points in cluster i. Since each point is a vector, this is a vector sum.
  • 1∣Ci∣\frac{1}{|C_i|}: We divide by the number of points to get the average.

For example, if cluster 1 contains three 2D points: (1,2)(1,2), (3,4)(3,4), and (5,6)(5,6), the centroid is:

μ1=13[(1,2)+(3,4)+(5,6)]=13(9,12)=(3,4)\mu_1 = \frac{1}{3}[(1,2) + (3,4) + (5,6)] = \frac{1}{3}(9, 12) = (3, 4)

The centroid (3,4)(3, 4) is the average position of the three points: it's the natural "center of mass" of the cluster.

The K-means Algorithm: An Iterative Solution

We now have all the pieces we need, but there's a challenge: finding the optimal clusters and centroids simultaneously is computationally intractable for most datasets. The problem is NP-hard, meaning there's no known efficient algorithm to find the globally optimal solution.

K-means solves this by using an iterative approach that alternates between two steps:

  1. Fix the centroids, optimize the cluster assignments: Given current centroid positions, assign each point to its nearest centroid.
  2. Fix the cluster assignments, optimize the centroids: Given current cluster assignments, recompute centroids as the mean of their assigned points.

By alternating between these two steps, the algorithm is guaranteed to converge to a local minimum of the objective function. While this might not be the global optimum, it's often a good solution in practice, especially with smart initialization strategies like k-means++.

Let's formalize this iterative process:

Step 1: Initialization

Choose k initial centroids μ1(0),μ2(0),…,μk(0)\mu_1^{(0)}, \mu_2^{(0)}, \ldots, \mu_k^{(0)}. The superscript (0)(0) indicates these are the centroids at iteration 0 (the starting point). The choice of initial centroids can significantly impact the final result, which is why k-means++ initialization is preferred over random selection.

Step 2: Assignment Step

For each data point xx, assign it to the cluster whose centroid is closest. Mathematically, at iteration tt, cluster ii contains all points for which centroid μi(t)\mu_i^{(t)} is the nearest:

Ci(t)={x:∥x−μi(t)∥2≤∥x−μj(t)∥2for all j≠i}C_i^{(t)} = \left\{ x : \|x - \mu_i^{(t)}\|^2 \leq \|x - \mu_j^{(t)}\|^2 \quad \text{for all } j \neq i \right\}

Let's decode this notation:

  • Ci(t)C_i^{(t)}: The set of points assigned to cluster i at iteration t.
  • xx: A data point we're considering.
  • ∥x−μi(t)∥2\|x - \mu_i^{(t)}\|^2: The squared distance from point x to centroid i.
  • "≤∥x−μj(t)∥2\leq \|x - \mu_j^{(t)}\|^2 for all j≠ij \neq i": Point x is closer to centroid i than to any other centroid j.

In plain English: "Cluster i contains all points that are closer to centroid i than to any other centroid."

This assignment step creates a Voronoi diagram in the feature space: each cluster's region consists of all points closer to its centroid than to any other centroid. The boundaries between clusters are hyperplanes that bisect the lines connecting adjacent centroids.

Let's visualize what a Voronoi diagram looks like:

Out[4]:
Visualization
Voronoi diagram showing four colored regions divided by centroids, with decision boundaries as perpendicular bisectors.
Voronoi diagram created by K-means cluster assignment. This visualization shows how K-means partitions the feature space into regions based on proximity to centroids. The four centroids (marked with red ''x'' symbols) divide the space into four colored regions. Any point within a colored region is closer to that region''s centroid than to any other centroid. The black lines are the decision boundaries: they are perpendicular bisectors of the lines connecting adjacent centroids. This geometric structure explains why K-means creates convex cluster regions and works best with spherical, well-separated clusters. Points near the boundaries are roughly equidistant from multiple centroids and could potentially be reassigned in future iterations.

Step 3: Update Step

Once we've assigned all points to clusters, we recompute each centroid as the mean of its assigned points:

μi(t+1)=1∣Ci(t)∣∑x∈Ci(t)x\mu_i^{(t+1)} = \frac{1}{|C_i^{(t)}|} \sum_{x \in C_i^{(t)}} x

The superscript (t+1)(t+1) indicates this is the new centroid for the next iteration. This update moves each centroid to the center of its cluster, which will generally reduce the within-cluster sum of squares.

Step 4: Convergence Check

We repeat steps 2 and 3 until one of these conditions is met:

  • The centroids no longer change (or change by less than a small threshold)
  • The cluster assignments no longer change
  • We reach a maximum number of iterations

The algorithm is guaranteed to converge because each step either decreases or maintains the objective function value (it can never increase). Eventually, we reach a point where no further improvement is possible, and the algorithm stops.

Compact Matrix Notation

For those familiar with linear algebra, we can express the K-means assignment step more compactly using matrix notation. Let X be our n×dn \times d data matrix (n points, d features) and M be our k×dk \times d centroid matrix (k centroids, d features).

For each point, the assignment step finds:

arg⁡min⁡i∥Xj−Mi∥2\arg\min_{i} \|X_j - M_i\|^2

where XjX_j is the j-th row of X (the j-th data point) and MiM_i is the i-th row of M (the i-th centroid). This notation emphasizes that we're finding the centroid index i that minimizes the distance for each point.

Mathematical Properties and Guarantees

Understanding the mathematical properties of K-means helps us know what to expect and when to trust the results:

Convergence Guarantee

K-means is guaranteed to converge to a local minimum of the objective function. This is because:

  1. The assignment step assigns each point to its nearest centroid, which can only decrease (or maintain) the total squared distance.
  2. The update step moves each centroid to the mean of its cluster, which minimizes the sum of squared distances for that cluster.
  3. Since both steps decrease or maintain the objective function, and there are only finitely many possible cluster assignments, the algorithm must eventually stop improving and converge.

However, this convergence is to a local minimum, not necessarily the global minimum. Different initializations can lead to different final clusterings, which is why running K-means multiple times with different random seeds (controlled by the n_init parameter) is recommended.

Computational Complexity

The algorithm has a time complexity of O(nkdt), where:

  • n: Number of data points
  • k: Number of clusters
  • d: Number of features (dimensions)
  • t: Number of iterations until convergence

In practice, t is typically much smaller than n (often less than 100 iterations), making K-means quite efficient even for large datasets. The algorithm scales linearly with the number of data points, which is excellent for big data applications.

Optimality of Centroids

The centroids produced by K-means are optimal in a specific sense: given a fixed cluster assignment, the centroids that minimize the within-cluster sum of squares are exactly the means of the clusters. This is a mathematical fact that follows from calculus: the derivative of the sum of squared distances with respect to the centroid position is zero at the mean.

This optimality property is what makes the update step so effective: by moving centroids to the mean of their clusters, we're guaranteed to improve (or at least not worsen) the objective function for that fixed assignment.

Space Partitioning

K-means creates a partition of the feature space into k regions, where each region contains all points closest to a particular centroid. These regions are convex polytopes (generalized polygons in higher dimensions) bounded by hyperplanes. This geometric structure explains why K-means works well for spherical, well-separated clusters but struggles with elongated or irregularly shaped clusters: the algorithm is fundamentally constrained to create convex regions.