From-scratch guide / Chapter 8 / Introductory
K-Means Clustering from Scratch in Python
K-means partitions samples by alternating two operations: assign each sample to its nearest centroid, then replace each centroid with the mean of its assigned samples. The loop below implements both operations and stops when the centroids no longer move materially.
How it works
minimize sum_i ||x_i - mu_(assignment_i)||^2Each sample x_i contributes its squared distance to the centroid mu of its assigned cluster. Assignment and update steps never increase this objective, but they can settle at a local minimum.
- 1Choose k initial centroids, using explicit or seeded initialization for reproducibility.
- 2Assign every sample to the nearest centroid by squared distance.
- 3Replace each non-empty centroid with the mean of its assigned samples.
- 4Repeat until centroid movement is below a tolerance or the iteration limit is reached.
Visual intuition
NumPy implementation
This example is deterministic and uses NumPy for arrays and arithmetic, not a prebuilt optimizer or model implementation.
import numpy as np
def assign_clusters(X, centroids):
squared_distances = ((X[:, None, :] - centroids[None, :, :]) ** 2).sum(axis=2)
return squared_distances.argmin(axis=1)
def update_centroids(X, labels, centroids):
updated = centroids.copy()
for cluster in range(len(centroids)):
members = X[labels == cluster]
if len(members):
updated[cluster] = members.mean(axis=0)
return updated
def kmeans(X, initial_centroids, max_iterations=100, tolerance=1e-8):
X = np.asarray(X, dtype=float)
centroids = np.asarray(initial_centroids, dtype=float).copy()
for _ in range(max_iterations):
labels = assign_clusters(X, centroids)
updated = update_centroids(X, labels, centroids)
if np.linalg.norm(updated - centroids) <= tolerance:
centroids = updated
break
centroids = updated
return centroids, assign_clusters(X, centroids)
X = np.array([[0, 0], [0, 1], [1, 0], [8, 8], [8, 9], [9, 8]])
centroids, labels = kmeans(X, initial_centroids=np.array([[0, 0], [9, 9]]))Broadcasting constructs every sample-to-centroid squared distance without a Python loop over samples.
An empty cluster keeps its previous centroid, making the edge-case policy explicit rather than producing a NaN mean.
The convergence check compares successive centroid arrays; assignments are recomputed once for the returned centroids.
What can go wrong
Initialization changes the answer
K-means can stop at different local minima, so practical runs compare multiple seeds or use a stronger initializer.
You still have to choose k
The algorithm optimizes for a supplied cluster count; domain evidence or validation must justify that count.
Distance reflects feature scale
A large-scale feature can dominate Euclidean distance unless the chosen representation or scaling is appropriate.
Clusters can become empty
This example retains the previous centroid; other explicit policies reinitialize it, but silently averaging zero members is invalid.
Continue with the book and source
This guide is grounded in Chapter 8, Linear Unsupervised Learning, and the official Machine Learning Refined notebooks. Use the chapter for context, then open the exact sources below for the full treatment and exercises.