k-Means Clustering

k-means clustering divides a set of points into $k$ groups so that each point is close to the mean of its group. It alternates two steps: assign every point to the nearest of $k$ centres, then move every centre to the mean of the points assigned to it. The iteration, usually called Lloyd’s algorithm, lowers the total squared distance from points to their centres at every step and stops at a local minimum of it. It is one of the simplest machine learning methods that need no labels; Elkan (2003) calls it by far the most widely used method for discovering clusters in data.

In computer vision the points are usually pixel colours, pixel colours with their positions, or local descriptors, and the centres become a palette, a set of regions or a vocabulary. k-means reduces an image to $k$ representative colours (colour quantization), segments it by colour or by colour and position, underlies the SLIC algorithm for superpixels, and builds the visual vocabularies of the bag of visual words. It measures similarity by the squared Euclidean distance, and that choice decides which clusters it can find.

Problem

Given $n$ points and a number $k$, find $k$ centres such that the sum of squared distances from each point to its nearest centre is as small as possible. Each centre then stands for a group of points, its cluster, and the nearest-centre rule partitions the whole space into $k$ cells. The same problem is called vector quantization in signal compression, where the centres are code vectors and the partition is the encoder (Bishop, 2006, §9.1.1; Hastie et al., 2009, §14.3.9).

Inputs and Outputs

• Input: a data matrix $X \in \mathbb{R}^{n \times d}$ with one point per row (scikit-learn’s layout), for example the $n = HW$ pixels of an $H \times W$ RGB image with $d = 3$; the number of clusters $k$; an initialization method; and a stopping rule.
• Output: $k$ centres $\boldsymbol{\mu}_1, \dots, \boldsymbol{\mu}_k \in \mathbb{R}^d$, a label $c_i \in \{1, \dots, k\}$ for every point, and the value of the objective, which scikit-learn calls the inertia. New points are labelled by their nearest centre.

Intuition

Suppose the centres were known: then each point should obviously go to its nearest centre. Suppose instead the groups were known: then the best single point to represent each group, in the sense of squared distance, is the group’s mean. k-means knows neither, so it guesses the centres, derives the groups, re-derives the centres from the groups, and repeats. Each half-step can only reduce the total squared distance, so the process settles, but where it settles depends on the first guess.

Squared Euclidean distance favours a particular kind of cluster. The boundary between two centres is the hyperplane halfway between them, whatever the spread or number of points on either side, so k-means works best when clusters are compact, roughly round, of similar spread, and well separated, and it carves elongated or very unequal groups into pieces of similar size instead.

Algorithm

Lloyd’s algorithm, as described by Bishop (2006, §9.1) and Hastie et al. (2009, Algorithm 14.1):

1. Initialize $k$ centres, for example $k$ distinct data points chosen at random, or by k-means++ seeding (see Parameters).
2. Assignment step. Give every point the label of its nearest centre, $c_i = \arg\min_j \lVert \mathbf{x}_i - \boldsymbol{\mu}_j \rVert^2$, breaking ties consistently.
3. Update step. Move every centre to the mean of the points assigned to it, $\boldsymbol{\mu}_j = \frac{1}{n_j} \sum_{i : c_i = j} \mathbf{x}_i$, where $n_j$ is the number of such points.
4. Repeat steps 2 and 3 until the assignments no longer change, or until the centres move less than a tolerance, or a maximum number of iterations is reached.

If a centre loses all its points, the update step is undefined. Implementations handle the empty cluster in different ways; scikit-learn 1.7.2 moves it to the point that lies farthest from its current centre and removes that point from its old cluster (our reading of its _relocate_empty_clusters_dense routine), and the code below uses a simplified form of the same rule, which leaves the point in its old cluster’s mean.

Bishop calls the assignment and update steps the E step and M step, because they are the hard-assignment limit of the expectation–maximization algorithm for a Gaussian mixture (see Mathematical Formulation).

Mathematical Formulation

The within-cluster sum of squares

With centres $\boldsymbol{\mu}_1, \dots, \boldsymbol{\mu}_k$ and labels $c_1, \dots, c_n$, k-means minimizes

$$J = \sum_{i=1}^{n} \lVert \mathbf{x}_i - \boldsymbol{\mu}_{c_i} \rVert^2 = \sum_{j=1}^{k} \sum_{i : c_i = j} \lVert \mathbf{x}_i - \boldsymbol{\mu}_j \rVert^2 ,$$

the within-cluster sum of squares (WCSS), which Bishop (§9.1, Eq. 9.1) writes with indicator variables $r_{ij} \in \{0, 1\}$ and calls a distortion measure. Minimizing over the labels for fixed centres gives the nearest-centre rule, so $J$ can also be written as a function of the centres alone, $J(\boldsymbol{\mu}_1, \dots, \boldsymbol{\mu}_k) = \sum_i \min_j \lVert \mathbf{x}_i - \boldsymbol{\mu}_j \rVert^2$, the form Arthur and Vassilvitskii call the potential $\phi$. Minimizing over the centres for fixed labels gives the cluster means. Dividing $J$ by $n$ gives the mean squared error per point, which for a colour image is the mean squared error between the image and its quantized version.

At the means, a cluster’s sum of squares equals half its sum of squared pairwise distances divided by its size: $\sum_{i \in C_j} \lVert \mathbf{x}_i - \bar{\mathbf{x}}_j \rVert^2 = \frac{1}{2 n_j} \sum_{i, i' \in C_j} \lVert \mathbf{x}_i - \mathbf{x}_{i'} \rVert^2$ (a standard identity, which follows from expanding the squares). k-means therefore prefers clusters whose members are close to one another, not only close to a centre. Because the total scatter of the data is fixed, minimizing the scatter within clusters is the same as maximizing the scatter between them (Hastie et al., 2009, §14.3.5).

The cells of the nearest-centre partition are the Voronoi cells of the centres (Hastie et al., §14.3.6). Each pair of cells is separated by the perpendicular bisector of the two centres, so every cluster k-means returns is a convex region of feature space.

Why each step lowers J

The assignment step cannot increase $J$, since each term $\lVert \mathbf{x}_i - \boldsymbol{\mu}_{c_i} \rVert^2$ is replaced by the smallest available one. For the update step, Arthur and Vassilvitskii (Lemma 2.1) use the identity, for a set $S$ with mean $\bar{\mathbf{x}}_S$ and any point $\mathbf{z}$,

$$\sum_{\mathbf{x} \in S} \lVert \mathbf{x} - \mathbf{z} \rVert^2 = \sum_{\mathbf{x} \in S} \lVert \mathbf{x} - \bar{\mathbf{x}}_S \rVert^2 + |S| \, \lVert \bar{\mathbf{x}}_S - \mathbf{z} \rVert^2 ,$$

so moving a centre $\mathbf{z}$ to its cluster’s mean lowers that cluster’s contribution by $|S| \lVert \bar{\mathbf{x}}_S - \mathbf{z} \rVert^2$. Equivalently, $J$ is a quadratic in each $\boldsymbol{\mu}_j$ with its minimum at the mean (Bishop, Eq. 9.3–9.4).

Convergence and local minima

$J$ never increases, and with consistent tie-breaking no partition can repeat; since there are at most $k^n$ partitions of $n$ points, the algorithm stops after a finite number of iterations (Arthur & Vassilvitskii, §1). It stops at a partition that neither step can improve, which may be far from the global minimum: Arthur and Vassilvitskii note that the ratio of the reached $J$ to the optimal one can be unbounded even for fixed $n$ and $k$, and with high probability even when the centres are chosen uniformly at random from the data. Hastie et al. (§14.3.5) explain why exhaustive search is out of the question: the number of ways to assign $n$ points to $k$ clusters (a Stirling number of the second kind) is already 34,105 for $n = 10$, $k = 4$, and about $10^{10}$ for $n = 19$, $k = 4$. Arthur and Vassilvitskii and Sculley (2010) state that solving the problem exactly is NP-hard.

scikit-learn’s documentation calls the algorithm very fast in practice, gives the average complexity as $O(knT)$ for $T$ iterations, and gives a worst-case bound of $O(n^{k+2/d})$ for $d$ features, referring to Arthur and Vassilvitskii’s paper “How Slow Is the k-Means Method?”.

k-means as a limit of a Gaussian mixture

k-means is the limiting case of a Gaussian mixture model whose components share the covariance $\epsilon I$, with $\epsilon$ fixed and taken to zero: the posterior probability that a point belongs to a component tends to one for the nearest mean and to zero for all others, the EM update of each mean becomes the k-means update, and maximizing the expected complete-data log-likelihood becomes minimizing $J$ (Bishop, 2006, §9.3.2; Hastie et al., 2009, §14.3.7). The mixture article gives the derivation. Read the other way (our gloss), k-means behaves as if every cluster were an isotropic Gaussian with the same small variance, which is where its assumptions come from: round clusters, equal spread, and a boundary that ignores how many points each cluster holds.

Parameters

The number of clusters $k$. Sometimes the application fixes it: a palette of 16 colours, a vocabulary of 10,000 visual words, a budget of 500 superpixels. When $k$ is meant to reflect natural groups in the data, it has to be estimated, and $J$ alone cannot do it, because $J$ keeps falling as $k$ grows, down to zero at $k = n$. Hastie et al. (§14.3.11) point out that this happens even on held-out data, since many centres fill the space densely, so cross-validation, the usual tool for choosing model size, does not apply. Common heuristics:

• Elbow. Plot $J$ against $k$ and look for the bend: while $k$ is below the number of well-separated groups, each new centre separates two groups and $J$ drops sharply; beyond it, each new centre splits a compact group and $J$ drops little (Hastie et al., §14.3.11).
• Silhouette. For each point, let $a$ be its mean distance to the other points of its cluster and $b$ its mean distance to the points of the nearest other cluster; its silhouette is $(b - a)/\max(a, b)$, between $-1$ and $1$, and the mean silhouette over all points is highest for compact, well-separated clusters (scikit-learn silhouette_score documentation).
• Gap statistic. Tibshirani, Walther and Hastie compare $\log J_k$ with its expected value for data drawn uniformly over a rectangle containing the data, and choose the smallest $k$ whose gap is within one simulation standard deviation of the gap at $k + 1$: a way of locating the elbow automatically, which also works when the data form a single cluster (as Hastie et al., §14.3.11, describe it).

Initialization. Because the result depends on the starting centres, the choice matters:

• Random data points. Choose $k$ distinct points uniformly at random, as Bishop (§9.1) recommends over arbitrary starting values. A random start often places two centres in one true cluster and none in another; Arthur and Vassilvitskii (§6.3) observe that the algorithm then never splits the merged clusters apart.
• k-means++. Arthur and Vassilvitskii choose the first centre uniformly at random from the data, and each subsequent centre from the data with probability proportional to $D(\mathbf{x})^2$, the squared distance from $\mathbf{x}$ to the nearest centre chosen so far (“$D^2$ weighting”), then run Lloyd’s algorithm. Points far from all existing centres, which are likely to lie in an uncovered cluster, are favoured. Their Theorem 3.1 bounds the expected objective after seeding alone, before any Lloyd iteration, by $\mathbb{E}[J] \le 8(\ln k + 2)\, J_{\mathrm{OPT}}$, so k-means++ is $O(\log k)$-competitive with the optimal clustering; Lloyd’s iterations can only lower $J$ further. They also show that $D^2$ seeding is no better than $2 \ln k$-competitive on some inputs (Theorem 4.1), so the bound is tight up to a constant factor. In their experiments (one synthetic and three real data sets, $k = 10, 25, 50$, 20 trials per setting), k-means++ had a lower average and a lower minimum potential than randomly seeded k-means in all 12 settings and a shorter mean running time in 11 (their Tables 1–4).
• Restarts. Run the algorithm from several starting points and keep the solution with the smallest $J$ (Hastie et al., §14.3.6). MacQueen (1967, §3.6) already suggested repeated runs from different starting points.

Stopping rule. Stop when no label changes, which is exact, or when the centres move less than a tolerance. scikit-learn 1.7.2 does both: it stops when the labels are unchanged, or when the sum of squared centre shifts falls below tol times the mean variance of the features (our reading of the source).

The feature space. k-means sees only Euclidean distances, so the representation decides what counts as similar: which colour space, whether position is included and with what weight, and whether features are standardized. A feature with a larger numerical range dominates the distance. Hastie et al. (§14.3.6) note that a weighted Euclidean distance can be obtained by rescaling the variables.

scikit-learn 1.7.2 defaults. KMeans(n_clusters=8, init="k-means++", n_init="auto", max_iter=300, tol=1e-4, algorithm="lloyd"), as the docstring of the installed release states. With n_init="auto" the number of runs is 1 when init is "k-means++" or an array and 10 when it is "random" or a callable; the "auto" option was added in version 1.2 and became the default in 1.4. Its k-means++ is a greedy variant: at each seeding step it draws several candidates, $2 + \lfloor \ln k \rfloor$ by default, and keeps the one that lowers the objective most, which the kmeans_plusplus documentation says worked better empirically than the original one-candidate form; Arthur and Vassilvitskii (§7) had found the same in their experiments, though their bound does not cover the greedy choice. algorithm="elkan" selects Elkan’s accelerated variant (see Variants).

Complexity

The assignment step computes $nk$ squared distances in $d$ dimensions, $O(nkd)$, and the update step costs $O(nd)$, so a run of $T$ iterations costs $O(nkdT)$ time and $O(nd + kd)$ memory beyond the labels. k-means++ seeding costs $O(nkd)$, one distance computation per point for each new centre. With restarts the cost is multiplied by their number. The $nk$ term is what limits very large problems: the bag of visual words article describes how vocabularies of $10^5$ to $10^6$ words over tens of millions of descriptors made exact iterations impractical, and the tree-structured and approximate nearest-centre searches that replaced them.

Implementation

The examples cluster pixel colours of synthetic images whose true colours and regions are known exactly. They implement k-means++ and Lloyd’s algorithm in NumPy, check them against scikit-learn’s KMeans, and were run with NumPy 2.3.5, SciPy 1.16.3 and scikit-learn 1.7.2, single-threaded (OMP_NUM_THREADS=1). Pixel values are in $[0, 1]$; the noise is not clipped.

Colour quantization and the choice of k

The first block paints a $128 \times 128$ image with six known colours in shapes of very different sizes (from 437 to 9,582 pixels), adds Gaussian noise with standard deviation $\sigma = 0.05$ per channel, and clusters its 16,384 pixel colours. It compares the six centres with the true palette, then fits $k = 1, \dots, 10$ and prints the distortion per pixel and the mean silhouette. For every $k$ it also runs scikit-learn’s Lloyd algorithm from the same k-means++ starting centres as the NumPy code and checks that labels and centres agree, and it prints the best of 10 runs from each implementation.

import numpy as np
from scipy.optimize import linear_sum_assignment

def sq_dists(X, C):
    """Squared Euclidean distances between the rows of X (n x d) and of C (k x d)."""
    return np.maximum((X**2).sum(1)[:, None] - 2 * X @ C.T + (C**2).sum(1)[None, :], 0)

def kmeans_pp(X, k, rng):
    """k-means++ seeding: first centre uniform, then each new centre drawn with
    probability proportional to D(x)^2, the squared distance to the nearest chosen centre."""
    C = [X[rng.integers(len(X))]]
    D2 = sq_dists(X, C[0][None]).ravel()
    for _ in range(k - 1):
        C.append(X[rng.choice(len(X), p=D2 / D2.sum())])
        D2 = np.minimum(D2, sq_dists(X, C[-1][None]).ravel())
    return np.array(C)

def lloyd(X, C, max_iter=300):
    """Lloyd's algorithm from initial centres C. Returns centres, labels and the
    objective J (sum of squared distances to the assigned centre) after each assignment."""
    C, labels, J = C.copy(), None, []
    for _ in range(max_iter):
        D = sq_dists(X, C)
        new = D.argmin(1)                                    # assignment step
        J.append(D[np.arange(len(X)), new].sum())
        if labels is not None and np.array_equal(new, labels):
            break
        labels = new
        for j in range(len(C)):                              # update step
            members = X[labels == j]
            if len(members):
                C[j] = members.mean(0)
            else:                                            # empty cluster: move it to the
                far = D[np.arange(len(X)), labels].argmax()  # point farthest from its centre
                C[j] = X[far]
    return C, labels, np.array(J)

def kmeans(X, k, n_init, rng):
    """Best of n_init runs of Lloyd's algorithm from k-means++ seeds, by final J."""
    return min((lloyd(X, kmeans_pp(X, k, rng)) for _ in range(n_init)), key=lambda r: r[2][-1])

# A 128 x 128 RGB image (values in [0, 1]) painted with six known colours, plus noise.
palette = np.array([[0.85, 0.80, 0.70],   # beige background
                    [0.80, 0.15, 0.15],   # red disk
                    [0.20, 0.60, 0.25],   # green square
                    [0.15, 0.25, 0.75],   # blue stripe
                    [0.95, 0.85, 0.20],   # yellow triangle
                    [0.15, 0.15, 0.20]])  # dark disk
H = W = 128
yy, xx = np.mgrid[0:H, 0:W]
truth = np.zeros((H, W), int)
truth[(yy - 40)**2 + (xx - 40)**2 < 22**2] = 1
truth[70:115, 15:60] = 2
truth[:, 100:112] = 3
truth[(yy > 60) & (yy < 120) & (xx > 65) & (xx < 95) & (yy - 60 > 2 * np.abs(xx - 80))] = 4
truth[(yy - 25)**2 + (xx - 80)**2 < 12**2] = 5
sigma = 0.05
rng = np.random.default_rng(0)
img = palette[truth] + rng.normal(0, sigma, (H, W, 3))
X = img.reshape(-1, 3)                                       # one pixel per row
print("pixels per colour:", np.bincount(truth.ravel()))

# Quantize to k = 6 colours and compare the centres with the palette.
C, labels, J = kmeans(X, 6, n_init=10, rng=rng)
print("J after each assignment step:", np.round(J, 2))
r, c = linear_sum_assignment(sq_dists(palette, C))           # match centres to palette
print("largest |centre - palette| over all channels: %.4f" % np.abs(C[c] - palette[r]).max())
agree = np.mean(c[truth.ravel()] == labels)
print("pixels whose cluster is their true colour's cluster: %.4f" % agree)
print("mean squared error per pixel: %.5f (noise alone: 3 sigma^2 = %.5f)"
      % (J[-1] / len(X), 3 * sigma**2))

# Distortion and silhouette as k grows; cross-check every k against scikit-learn.
import sklearn
from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score
print(f"\nscikit-learn {sklearn.__version__}")
print(" k   J/n (ours)  J/n (sklearn)  silhouette  same as sklearn from same seeds")
for k in range(1, 11):
    C0 = kmeans_pp(X, k, rng)
    Ck, lk, Jk = lloyd(X, C0)
    km = KMeans(k, init=C0, n_init=1, tol=0, algorithm="lloyd").fit(X)
    same = np.array_equal(lk, km.labels_) and np.allclose(Ck, km.cluster_centers_)
    best = kmeans(X, k, n_init=10, rng=rng)
    sk = KMeans(k, n_init=10, random_state=0).fit(X).inertia_
    sil = silhouette_score(X, best[1], sample_size=4000, random_state=0) if k > 1 else np.nan
    print(f"{k:2d}   {best[2][-1] / len(X):.5f}     {sk / len(X):.5f}        {sil:6.3f}      {same}")

Output:

pixels per colour: [9582 1513 2025 1536 1291  437]
J after each assignment step: [236.87 123.21 123.2 ]
largest |centre - palette| over all channels: 0.0026
pixels whose cluster is their true colour's cluster: 1.0000
mean squared error per pixel: 0.00752 (noise alone: 3 sigma^2 = 0.00750)

scikit-learn 1.7.2
 k   J/n (ours)  J/n (sklearn)  silhouette  same as sklearn from same seeds
 1   0.21076     0.21076           nan      True
 2   0.08599     0.08599         0.642      True
 3   0.04948     0.04948         0.705      True
 4   0.03038     0.03038         0.752      True
 5   0.01204     0.01204         0.768      True
 6   0.00752     0.00752         0.784      True
 7   0.00656     0.00657         0.452      True
 8   0.00592     0.00591         0.450      True
 9   0.00578     0.00544         0.410      True
10   0.00529     0.00529         0.402      True

The noise is small compared with the distances between the palette colours, so with $k = 6$ every pixel lands in its colour’s cluster and the centres match the palette to within 0.003 in every channel, about the standard error of a mean of a few hundred pixels with $\sigma = 0.05$. The remaining distortion per pixel equals $3\sigma^2$, the noise variance summed over three channels: once the colours are separated, what is left is noise, which no choice of six centres can remove. Started from k-means++ seeds, the best run converged after two updates.

The distortion curve has its elbow at $k = 6$: each of the first five extra centres separates colours that were merged and removes a large share of $J$, while centres beyond the sixth only split a noise cloud, which removes a small fraction of its variance. The silhouette peaks at the same $k$ and falls sharply at $k = 7$, because a split noise cloud gives two clusters whose points are nearly as close to the other half as to their own. From identical starting centres the NumPy and scikit-learn implementations agree exactly for every $k$. Their best-of-10 values agree to within 0.2% except at $k = 9$, where the two sets of runs ended in different local minima, one about 6% higher than the other: past the natural number of clusters, the extra centres can split the noise clouds in many nearly equivalent ways.

Initialization and restarts

The second block uses a colour chart of 12 equal patches of $32 \times 32$ pixels with known colours and noise $\sigma = 0.04$, and compares random-point seeding with k-means++ over 50 seeds, one run each, then with the best of 10 runs over 20 seeds, and finally scikit-learn’s default (greedy k-means++, one run). A colour counts as found when some centre lies within 0.03 of it in every channel; $J_{\mathrm{ref}}$ is the objective after running Lloyd’s algorithm from the true colours.

import numpy as np
# Uses sq_dists, kmeans_pp and lloyd from the first block.

# A synthetic colour chart: 12 known colours in a 3 x 4 grid of 32 x 32 patches, plus noise.
palette = np.array([[0.9, 0.9, 0.9], [0.1, 0.1, 0.1], [0.8, 0.2, 0.2], [0.2, 0.7, 0.2],
                    [0.2, 0.3, 0.8], [0.9, 0.8, 0.2], [0.8, 0.3, 0.8], [0.2, 0.8, 0.8],
                    [0.5, 0.5, 0.5], [0.9, 0.5, 0.1], [0.4, 0.2, 0.1], [0.5, 0.8, 0.5]])
truth = np.kron(np.arange(12).reshape(3, 4), np.ones((32, 32), int))
H, W = truth.shape
rng = np.random.default_rng(1)
X = (palette[truth] + rng.normal(0, 0.04, (H, W, 3))).reshape(-1, 3)

k = 12
J_ref = lloyd(X, palette)[2][-1]                             # Lloyd started at the true colours

def found(C):
    """Number of true colours that have a centre within 0.03 in every channel."""
    return int((np.abs(palette[:, None] - C[None]).max(2) < 0.03).any(1).sum())

def random_init(X, k, rng):
    return X[rng.choice(len(X), k, replace=False)]           # k distinct random pixels

print(f"J of the solution started at the true colours: {J_ref:.2f}")
print("one run per seed, seeds 0-49:")
for name, init in (("random pixels", random_init), ("k-means++", kmeans_pp)):
    ratio, nfound, iters = [], [], []
    for seed in range(50):
        C, lab, J = lloyd(X, init(X, k, np.random.default_rng(seed)))
        ratio.append(J[-1] / J_ref); nfound.append(found(C)); iters.append(len(J))
    ratio, nfound = np.array(ratio), np.array(nfound)
    print(f"  {name:13s}: J/J_ref min {ratio.min():.3f}, median {np.median(ratio):.3f}, "
          f"max {ratio.max():.3f}; all 12 colours found in {np.mean(nfound == 12):.0%} "
          f"of runs (fewest {nfound.min()}); iterations median {int(np.median(iters))}, "
          f"max {max(iters)}")

print("best of 10 runs by J, seeds 0-19:")
for name, init in (("random pixels", random_init), ("k-means++", kmeans_pp)):
    nfound = []
    for seed in range(20):
        rng = np.random.default_rng(seed)
        C = min((lloyd(X, init(X, k, rng)) for _ in range(10)), key=lambda r: r[2][-1])[0]
        nfound.append(found(C))
    print(f"  {name:13s}: all 12 colours found for {np.mean(np.array(nfound) == 12):.0%} of seeds")

# scikit-learn's default: greedy k-means++ (2 + int(ln k) candidates per centre), one run.
from sklearn.cluster import KMeans
nf = [found(KMeans(k, random_state=s).fit(X).cluster_centers_) for s in range(50)]
print(f"scikit-learn KMeans defaults, seeds 0-49: all 12 colours found in "
      f"{np.mean(np.array(nf) == 12):.0%} of runs (fewest {min(nf)})")

# One bad random start in detail: J after each assignment step never increases.
for seed in range(50):
    C, lab, J = lloyd(X, random_init(X, k, np.random.default_rng(seed)))
    if found(C) < 12:
        print(f"seed {seed}, random start: J =", np.round(J[:8], 1), "...", np.round(J[-1], 2),
              f"after {len(J)} steps; colours found {found(C)}; monotone: {np.all(np.diff(J) <= 1e-9)}")
        break

Output:

J of the solution started at the true colours: 58.04
one run per seed, seeds 0-49:
  random pixels: J/J_ref min 1.000, median 3.901, max 8.174; all 12 colours found in 2% of runs (fewest 3); iterations median 27, max 70
  k-means++    : J/J_ref min 1.000, median 1.770, max 5.016; all 12 colours found in 28% of runs (fewest 7); iterations median 14, max 67
best of 10 runs by J, seeds 0-19:
  random pixels: all 12 colours found for 25% of seeds
  k-means++    : all 12 colours found for 100% of seeds
scikit-learn KMeans defaults, seeds 0-49: all 12 colours found in 88% of runs (fewest 10)
seed 0, random start: J = [760.9 530.9 393.4 299.1 201.2 166.1 156.3 124. ] ... 102.61 after 25 steps; colours found 10; monotone: True

Most single runs end in a local minimum in which some patches share a centre and other patches are split between two. With random seeding only 1 run in 50 found all 12 colours, and the median run ended with nearly four times $J_{\mathrm{ref}}$. k-means++ seeding found all 12 in 28% of single runs, cut the median excess over $J_{\mathrm{ref}}$ by about three quarters, and needed about half as many iterations, since its seeds start in different patches more often; its worst run, at about five times $J_{\mathrm{ref}}$, is still far inside the $8(\ln 12 + 2) \approx 36$ factor of Arthur and Vassilvitskii’s bound, which in any case concerns the average. Keeping the best of 10 runs fixed k-means++ for all 20 seeds but random seeding for only 5 of them. scikit-learn’s greedy seeding succeeded in 88% of single runs. The last line follows one failed random start: $J$ falls at every step, as the theory requires, and stops at about 1.8 times $J_{\mathrm{ref}}$ with two colours missed.

Where k-means fails

The third block builds three images in which the true regions are known and k-means, run as above with 10 k-means++ restarts, does not recover them. In the first, a textured background with per-channel standard deviation 0.10 covers 90% of the image and a flat object with standard deviation 0.02 covers 10%. In the second, a background and a disk of two different paints are lit by the same illumination ramp, brightness rising from 0.25 to 1.0 across the image, so each region’s colours spread along a line in RGB space; the block clusters the RGB values and then the chromaticities $r = R/(R+G+B)$, $g = G/(R+G+B)$, which remove brightness. In the third, two red disks of the same colour on a beige background are clustered into 16 groups in a five-dimensional space of colour and position, with the position coordinates scaled to $[0, 1]$ and multiplied by a weight $w$.

import numpy as np
from scipy.optimize import linear_sum_assignment
# Uses kmeans (and the functions it calls) from the first block.

def matched_accuracy(truth, labels):
    """Fraction of pixels in the right cluster under the best one-to-one matching
    of true regions to clusters."""
    M = np.zeros((truth.max() + 1, labels.max() + 1))
    np.add.at(M, (truth, labels), 1)
    r, c = linear_sum_assignment(-M)
    return M[r, c].sum() / len(truth)

def wcss(X, labels):
    """Within-cluster sum of squares of a given partition."""
    return sum(((X[labels == j] - X[labels == j].mean(0))**2).sum() for j in np.unique(labels))

H = W = 128
yy, xx = np.mgrid[0:H, 0:W]
rng = np.random.default_rng(0)

# 1. Unequal spread: a strongly textured background (90% of pixels, per-channel std 0.10)
#    and a small flat object (std 0.02) whose colour differs by 0.4 in the red channel.
obj = ((yy - 64)**2 + (xx - 64)**2 < 23**2).ravel().astype(int)
noise = np.where(obj[:, None] == 1, 0.02, 0.10) * rng.normal(size=(H * W, 3))
X = np.where(obj[:, None] == 1, [0.8, 0.4, 0.4], [0.4, 0.4, 0.4]) + noise
C, lab, J = kmeans(X, 2, 10, rng)
o = lab[obj == 1].mean() > 0.5                                # cluster holding the object
print(f"1. object pixels: {obj.mean():.1%}; cluster labelled 'object' holds "
      f"{np.mean(lab == o):.1%} of the image, {np.mean(lab[obj == 0] == o):.1%} of the background")
print(f"   J of k-means {J[-1]:.1f} < J of the true partition {wcss(X, obj):.1f}")

# 2. Elongated cluster: a background of one paint under an illumination ramp (brightness
#    0.25 to 1.0 from left to right) and a disk of a second paint, lit the same way.
shade = (0.25 + 0.75 * xx / (W - 1)).ravel()[:, None]
disk = ((yy - 64)**2 + (xx - 90)**2 < 20**2).ravel().astype(int)
albedo = np.where(disk[:, None] == 1, [0.30, 0.60, 0.30], [0.60, 0.45, 0.35])
X = albedo * shade + rng.normal(0, 0.01, (H * W, 3))
rgb = kmeans(X, 2, 10, rng)[1]
chrom = X / X.sum(1, keepdims=True)                         # (r, g, b) / (r + g + b)
cro = kmeans(chrom[:, :2], 2, 10, rng)[1]
print(f"2. disk pixels: {disk.mean():.1%}; k = 2, matched accuracy {matched_accuracy(disk, rgb):.3f} "
      f"in RGB, {matched_accuracy(disk, cro):.3f} in chromaticity (r, g)")

# 3. Feature scaling: two red disks of the same colour on a beige background, clustered
#    into k = 16 groups in colour + position, position (x, y in [0, 1]) multiplied by w.
t = np.zeros((H, W), int)
t[(yy - 64)**2 + (xx - 32)**2 < 18**2] = 1
t[(yy - 64)**2 + (xx - 96)**2 < 18**2] = 2
t = t.ravel()
col = np.where(t[:, None] > 0, [0.8, 0.2, 0.2], [0.85, 0.8, 0.7]) + rng.normal(0, 0.05, (H * W, 3))
pos = np.stack([xx.ravel(), yy.ravel()], 1) / (W - 1)
for w in (0.0, 0.1, 0.3, 1.0, 3.0, 10.0):
    lab = kmeans(np.hstack([col, w * pos]), 16, 3, rng)[1]
    M = np.zeros((3, 16)); np.add.at(M, (t, lab), 1)
    mixed = 1 - M.max(0).sum() / len(t)                      # pixels outside their cluster's majority region
    both = int(((M[1] >= 20) & (M[2] >= 20)).sum())          # clusters spanning both disks
    print(f"3. w = {w:4.1f}: pixels not in their cluster's majority region {mixed:.3f}, "
          f"clusters spanning both disks {both}")

Output:

1. object pixels: 10.1%; cluster labelled 'object' holds 14.1% of the image, 4.4% of the background
   J of k-means 432.1 < J of the true partition 445.7
2. disk pixels: 7.6%; k = 2, matched accuracy 0.547 in RGB, 1.000 in chromaticity (r, g)
3. w =  0.0: pixels not in their cluster's majority region 0.060, clusters spanning both disks 4
3. w =  0.1: pixels not in their cluster's majority region 0.059, clusters spanning both disks 4
3. w =  0.3: pixels not in their cluster's majority region 0.000, clusters spanning both disks 0
3. w =  1.0: pixels not in their cluster's majority region 0.000, clusters spanning both disks 0
3. w =  3.0: pixels not in their cluster's majority region 0.000, clusters spanning both disks 0
3. w = 10.0: pixels not in their cluster's majority region 0.062, clusters spanning both disks 0

In the first image, k-means puts the boundary roughly halfway between the two colour means, so background pixels whose texture pushes them towards the object’s colour are labelled as object, and the object cluster is 40% larger than the object. This is not a failure of the optimizer: the wrong partition has a lower $J$ than the true one, so it is what the objective asks for when clusters differ in spread. A mixture model that estimates a variance per cluster would, by our reasoning, place the boundary much closer to the tight cluster.

In the second image the shading stretches each paint into a long segment in RGB space, longer than the distance between the two paints, so k-means with $k = 2$ cuts the image into a dark and a bright half; its matched accuracy is below the 0.924 that labelling every pixel as background would score. In chromaticity coordinates each paint collapses to a point and k-means separates them perfectly. The fix is a better feature space, not a better algorithm, and normalizing out brightness this way is a standard choice for colour segmentation, as the colour spaces article notes.

In the third image the weight on position decides the result. With $w = 0$ (and $w = 0.1$) clustering is by colour alone, and four of the 16 clusters contain pixels of both disks, so about 6% of the pixels fall outside their cluster’s majority region. With $w$ between 0.3 and 3, position separates the disks and the clusters are pure. With $w = 10$ position dominates, the clusters become spatial cells that ignore colour, and about 6% of the pixels again sit in mixed clusters, now ones that straddle a disk’s edge. SLIC’s compactness parameter $m$ makes the same trade-off for superpixels.

Properties and Behavior

• Monotone descent to a local minimum. $J$ never increases and the algorithm stops after finitely many iterations, at a partition that depends on the start; restarts and careful seeding are the practical remedies.
• Hard, convex, Voronoi partition. Every point belongs to exactly one cluster, and clusters are the Voronoi cells of the centres; points halfway between two centres get the same confident label as points at a centre, which the soft assignments of a Gaussian mixture avoid (Bishop, §9.1).
• Centres are means. They are generally not data points (the scikit-learn user guide notes this), and they are pulled by outliers: squared distance gives the largest distances the most influence, so k-means lacks robustness to outliers (Bishop, §9.1; Hastie et al., §14.3.10).
• Sensitivity to scale. Multiplying one feature by a constant changes the clustering, as the third example shows; since only Euclidean distances matter, it follows that k-means is unaffected by rotating or translating the feature space, or by scaling all features by the same factor, but not by scaling them separately.
• Reproducible only with a fixed seed. With random seeding or k-means++ two runs can give different clusterings; fixing the random state, or the initial centres, makes them repeatable.

Limitations

• $k$ must be given. The heuristics above estimate it but do not always agree, and many data sets have no single natural number of clusters.
• Spherical, similar-sized clusters. The equal-distance boundary assumes clusters of similar spread; clusters that are elongated, curved, of very different spread, or nested in one another are cut or merged, as the scikit-learn user guide also warns (it says that inertia assumes convex, isotropic clusters).
• Euclidean data only. The mean must be defined, so categorical data or data compared by an arbitrary dissimilarity need k-medoids or another method (Bishop, §9.1).
• High dimensions. In very high-dimensional spaces Euclidean distances tend to become inflated and less informative; the scikit-learn user guide suggests reducing dimension first, for example with principal component analysis.
• No spatial coherence in segmentation. Clustering pixel colours ignores where pixels are, so segments can be fragmented and noisy; Bishop (§9.1.1) uses k-means segmentation only as an illustration for this reason. Adding position, as SLIC does, or using methods built on pixel neighbourhoods such as normalized cuts, addresses it.

Variants

• Online (sequential) k-means. MacQueen’s original procedure processes points one at a time: each new point joins the group with the nearest mean, and that mean is updated to include it (MacQueen, 1967, §2.1). Bishop (§9.1) derives the general sequential update $\boldsymbol{\mu}_j \leftarrow \boldsymbol{\mu}_j + \eta_t (\mathbf{x}_t - \boldsymbol{\mu}_j)$ with a decreasing learning rate as a Robbins–Monro procedure, which is stochastic gradient descent on $J$.
• Mini-batch k-means. Sculley (2010) assigns a random mini-batch of points to their nearest centres and then moves each centre towards each of its batch points with a per-centre learning rate $1/v$, where $v$ counts the points the centre has received so far, so each centre is a running mean. On 781,265 training documents of the RCV1 collection, it reached near-optimal held-out objective values orders of magnitude faster than batch k-means, and better solutions than single-point stochastic gradient descent (his Fig. 1). scikit-learn’s MiniBatchKMeans (default batch_size=1024) implements it and its user guide describes the results as generally only slightly worse than the standard algorithm.
• Elkan’s accelerated k-means. Elkan (2003) avoids most distance computations while producing exactly the same centres as the standard algorithm after every iteration. If $d(\mathbf{b}, \mathbf{c}) \ge 2 d(\mathbf{x}, \mathbf{b})$ then $d(\mathbf{x}, \mathbf{c}) \ge d(\mathbf{x}, \mathbf{b})$ (his Lemma 1), so a point cannot switch to a centre that is far from its own; and $d(\mathbf{x}, \mathbf{c}) \ge \max\{0, d(\mathbf{x}, \mathbf{b}) - d(\mathbf{b}, \mathbf{c})\}$ (Lemma 2) lets lower bounds on point-to-centre distances be carried across iterations as centres move. Each point keeps an upper bound to its own centre and a lower bound to every other centre. In scikit-learn (algorithm="elkan") this costs an extra array of shape $n \times k$, and the documentation says it can be more efficient on data with well-defined clusters.
• k-medoids. Replaces squared Euclidean distance by an arbitrary dissimilarity and restricts each centre to be one of its cluster’s points, which makes it applicable to any dissimilarity, removes the outsized influence of large distances, and costs an update step quadratic in cluster size (Bishop, §9.1; Hastie et al., §14.3.10).
• Hierarchical (bisecting) k-means. Splitting the data with 2-means and recursing gives a tree of clusters, a divisive form of hierarchical clustering whose result depends on the starting configuration at each split (Hastie et al., §14.3.12); the vocabulary tree of the bag of visual words uses $b$-way splits for fast quantization into millions of words. scikit-learn provides BisectingKMeans.
• Gaussian mixtures. Soft assignments and a covariance per cluster remove the equal-spread, spherical assumption; Bishop (§9.1) notes that k-means is often used to initialize a mixture before EM, which is scikit-learn’s default for GaussianMixture.

Use in Computer Vision

Colour quantization and compression. Clustering the pixel colours of an image and replacing every pixel by its centre gives a $k$-colour palette image. Bishop (§9.1.1) counts the cost: $24N$ bits for $N$ pixels at 8 bits per channel, against $N \log_2 k$ bits for the labels plus $24k$ for the palette, about 4.2%, 8.3% and 16.7% of the original for $k = 2, 3, 10$ on his $240 \times 180$ example. Clustering blocks of pixels instead of single pixels is vector quantization proper: Hastie et al. (§14.3.9) quantize $2 \times 2$ blocks of a $1024 \times 1024$ grayscale photograph with $k = 200$ and $k = 4$ codewords, at 1.91 and 0.50 bits per pixel, and note that k-means is called Lloyd’s algorithm in this setting.

Segmentation. Clustering pixels by colour, or by colour and position, is the simplest image segmentation by clustering. In one dimension, the iterative threshold selection of Ridler and Calvard amounts to two-cluster k-means on intensities, as the thresholding article notes.

Superpixels. SLIC runs k-means in CIELAB colour and image position with a weighted distance, searching only a window around each centre, so each iteration costs time linear in the number of pixels rather than in pixels times centres; the superpixels article gives the details.

Visual vocabularies. The bag of visual words quantizes local descriptors such as SIFT against a k-means codebook and represents each image by its histogram of visual words; the bag of visual words article covers vocabulary size, the vocabulary tree and approximate k-means. Dense trajectory video descriptors were encoded the same way, with codebooks of 4,000 words per descriptor type (dense trajectories).

Search and compression of descriptors. Inverted-file indexes assign each database vector to the nearest of a set of k-means centroids, and product quantization compresses vectors with a separate k-means codebook per subvector (approximate nearest neighbor search).

Learning. YOLOv2 chose its anchor boxes by clustering the widths and heights of training boxes with k-means under the distance $1 - \mathrm{IoU}$ instead of the Euclidean one. DeepCluster trains a network by alternating k-means on its current features with learning to predict the cluster indices (self-supervised learning). k-means on embeddings is also used to group or deduplicate unlabelled image collections.

History

MacQueen (1967) named the method: his “k-means” process starts with $k$ groups of one random point each and adds each new point to the group with the nearest mean, updating that mean, and he studied its behaviour as the sample grows. He also described a two-step batch procedure, compute the means of the current groups and then regroup the points by nearest mean, which Edward Forgy and Robert Jennrich had independently suggested to him in 1963; it lowers the within-group variance at every iteration and converges in a finite number of steps. He named Sebestyen’s 1962 monograph as the earliest explicit use of the process he knew of. The batch iteration is today usually called Lloyd’s algorithm, after Stuart P. Lloyd’s paper on least-squares quantization for pulse-code modulation, published in 1982 (Bishop, 2006, Ch. 9; Arthur & Vassilvitskii, 2007). By its abstract, that paper derives necessary conditions for an optimal quantizer of signal amplitudes with a finite number of levels, minimizing the mean quantization noise power, which is the one-dimensional form of the same problem (our reading). Arthur and Vassilvitskii’s k-means++ seeding (SODA 2007, after a 2006 Stanford technical report) gave it a guarantee on the expected objective, which randomly seeded k-means lacks.