K-means Clustering › Formula
Machine Learning from ScratchAll chapters 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 . Each cluster is simply a set of data points that belong together. For each cluster, we'll compute a centroid 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:
Let's unpack what each part means:
- : The k clusters we're trying to find. Each cluster is a set of data points.
- : "For each data point x that belongs to cluster i." This notation tells us to consider every point in the cluster.
- : The centroid (mean position) of cluster i. This is the point that represents the cluster's center.
- : The squared Euclidean distance between point x and its cluster's centroid . This measures how far the point is from its cluster center.
- The double summation: The inner sum () adds up distances for all points in cluster i. The outer sum () adds up the results across all k clusters.
Our goal is to find the cluster assignments and centroids that minimize this total squared distance:
Why squared distances? You might wonder why we square the distances rather than just using regular distances. There are two important reasons:
-
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.
-
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:

Understanding Distance in Detail
Let's zoom in on the distance calculation itself. When we write , 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 is a d-dimensional vector.
The squared Euclidean distance is calculated by summing the squared differences across all dimensions:
The components in this expression have the following meanings:
- : The number of features (dimensions) in our data. For example, if we're clustering customers by age and income, then d = 2.
- : The value of the j-th feature for data point x. For instance, might be the income of a customer.
- : The value of the j-th feature for centroid . This is the average value of feature j across all points in cluster i.
- : 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 and a centroid , the squared distance is:
Let's visualize this distance calculation geometrically:

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:
Breaking this down:
- : The number of points in cluster i. The vertical bars denote the cardinality (size) of the set.
- : The sum of all data points in cluster i. Since each point is a vector, this is a vector sum.
- : We divide by the number of points to get the average.
For example, if cluster 1 contains three 2D points: , , and , the centroid is:
The centroid 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:
- Fix the centroids, optimize the cluster assignments: Given current centroid positions, assign each point to its nearest centroid.
- 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 . The superscript 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 , assign it to the cluster whose centroid is closest. Mathematically, at iteration , cluster contains all points for which centroid is the nearest:
Let's decode this notation:
- : The set of points assigned to cluster i at iteration t.
- : A data point we're considering.
- : The squared distance from point x to centroid i.
- " for all ": 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:

Step 3: Update Step
Once we've assigned all points to clusters, we recompute each centroid as the mean of its assigned points:
The superscript 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 data matrix (n points, d features) and M be our centroid matrix (k centroids, d features).
For each point, the assignment step finds:
where is the j-th row of X (the j-th data point) and 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:
- The assignment step assigns each point to its nearest centroid, which can only decrease (or maintain) the total squared distance.
- The update step moves each centroid to the mean of its cluster, which minimizes the sum of squared distances for that cluster.
- 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.