Introduction

A solid K-means implementation has two jobs: choose initial centroids and then alternate between assignment and update steps until the centroids stop moving. The algorithm is simple enough to write from scratch, but it is also easy to implement badly if you ignore empty clusters, initialization quality, or convergence checks.

Core Idea of K-Means

Given k clusters and a set of numeric points, K-means repeatedly:

1. assigns each point to its nearest centroid
2. recomputes each centroid as the mean of the points assigned to it
3. stops when the centroids barely change or after a maximum number of iterations

Vanilla random initialization works, but k-means++ usually converges to a better solution because it spreads the initial centroids more intelligently.

A Runnable NumPy Implementation

The example below uses k-means++ style initialization and then runs the normal K-means loop.

1import numpy as np
2
3
4def init_kmeans_pp(X, k, rng):
5    centroids = [X[rng.integers(len(X))]]
6    for _ in range(1, k):
7        distances = np.min(
8            np.sum((X[:, None, :] - np.array(centroids)[None, :, :]) ** 2, axis=2),
9            axis=1,
10        )
11        probs = distances / distances.sum()
12        index = rng.choice(len(X), p=probs)
13        centroids.append(X[index])
14    return np.array(centroids, dtype=float)
15
16
17def kmeans(X, k, max_iter=100, tol=1e-4, seed=0):
18    rng = np.random.default_rng(seed)
19    centroids = init_kmeans_pp(X, k, rng)
20
21    for _ in range(max_iter):
22        distances = np.sum((X[:, None, :] - centroids[None, :, :]) ** 2, axis=2)
23        labels = np.argmin(distances, axis=1)
24
25        new_centroids = centroids.copy()
26        for i in range(k):
27            members = X[labels == i]
28            if len(members) > 0:
29                new_centroids[i] = members.mean(axis=0)
30
31        shift = np.max(np.linalg.norm(new_centroids - centroids, axis=1))
32        centroids = new_centroids
33        if shift < tol:
34            break
35
36    return centroids, labels
37
38X = np.array([
39    [1.0, 1.0], [1.2, 0.8], [0.8, 1.1],
40    [4.0, 4.0], [4.2, 3.8], [3.9, 4.1],
41])
42
43centroids, labels = kmeans(X, k=2)
44print(centroids)
45print(labels)

This is enough to experiment with the algorithm and inspect how the centroids move.

Why Initialization Matters

K-means optimizes a non-convex objective, so different starting centroids can lead to different final clusters. Random starts may work on easy datasets, but they often converge to poor local optima.

k-means++ improves this by preferring new centroids that are far from existing ones. That does not guarantee the global optimum, but it usually reduces bad initializations and stabilizes results.

In practice, many production implementations still run multiple initializations and keep the best solution by inertia.

Handling Empty Clusters

One practical issue is empty clusters. During an iteration, a centroid may end up with no assigned points. If you blindly compute its mean, the centroid becomes invalid.

Common fixes are:

• keep the previous centroid unchanged
• reinitialize it to a random point
• reinitialize it to the point with the largest current error

The example above keeps the old centroid when a cluster becomes empty. That is simple and stable enough for a teaching implementation.

Choosing k

K-means requires you to choose the number of clusters in advance. The algorithm will always return exactly k clusters even if the data does not naturally support that choice.

Typical heuristics include:

• elbow plots of inertia
• silhouette scores
• domain-specific expectations about the data

Those methods guide the choice, but they do not remove the need for judgment.

Common Pitfalls

A common mistake is using K-means on features with very different scales. Since the algorithm relies on Euclidean distance, large-scale features dominate unless you standardize first.

Another issue is expecting K-means to find non-spherical or strongly unbalanced clusters. The algorithm works best when clusters are roughly compact and mean-based.

A third problem is relying on one random initialization and assuming the result is stable.

Summary

• K-means alternates between assigning points and recomputing centroids.
• 'k-means++ gives a stronger initialization than plain random seeding.'
• A good implementation needs convergence checks and empty-cluster handling.
• Standardize features when scale differences matter.
• Validate both the chosen k and the stability of the final clustering.