Writing K-Means from Scratch: One Developer's Journey into Machine Learning

Why It's Worth Committing to Implementing Algorithms from Scratch

In an era where the machine learning toolchain is highly mature, calling sklearn.cluster.KMeans takes just one line of code. But a developer on Reddit chose a harder path: writing every algorithm from scratch. According to their share, this is already day 59 of the "building machine learning algorithms from scratch" series — they've just finished every algorithm in the supervised learning section and have implemented K-Means clustering from unsupervised learning.

The value of this exercise lies in penetrating understanding of the algorithm's internal mechanics. K-Means seems simple — randomly pick centers, assign samples, update centers, iterate to convergence — but when you implement it by hand you run into a series of engineering details: how should cluster centers be initialized to avoid local optima? Should the distance metric be Euclidean distance or cosine similarity? Should the iteration stopping criterion be a threshold on center displacement or on the change in the objective function? The answers to these questions can only be truly felt by someone who has written the code, run it on data, and tuned the parameters.

The Core Logic of K-Means and Key Implementation Points

K-Means aims to partition n samples into k clusters so that the sum of squared distances from samples to their cluster centers is minimized. When implementing it by hand, the following aspects deserve special attention:

Initialization Strategy Determines the Lower Bound

Naive random initialization easily falls into poor local optima. In practice, the K-Means++ strategy is commonly used: the first center is chosen randomly, and each subsequent center is chosen with probability proportional to the squared distance from existing centers, so that initial centers are as spread out as possible. If you skip this step when writing by hand, you may not notice a difference on small datasets, but on data with complex cluster structure the results will fluctuate noticeably more.

The probabilistic selection mechanism of K-Means++: Let the set of existing centers be C. For each sample point x, compute its distance D(x) to the nearest center. The probability that sample x is chosen as the new center is proportional to D(x)², i.e., P(x) = D(x)² / Σ D(x')². This squared weighting gives points farther away a higher chance of being selected, thereby pushing initial centers away from each other. Although this strategy cannot guarantee global optimality, it improves the expected approximation ratio from arbitrarily bad to O(log k), making it the default initialization scheme in practice.

Distance Computation and Vectorization

Computing distances with a per-sample loop is extremely inefficient in Python. Even with a pure NumPy implementation, the distance matrix should be written in broadcast form as (X[:, np.newaxis] - centers) ** 2 to avoid explicit for loops. The degree of vectorization here directly determines the algorithm's usability on medium-scale data.

Convergence Criteria and Empty Cluster Handling

The iteration termination condition is usually set as center displacement below a threshold, or cluster assignments unchanged across two consecutive rounds. Another easily overlooked edge case is the empty cluster: when a center is not assigned any samples, that center needs to be reinitialized (e.g., pick the sample point farthest from the existing centers), otherwise the algorithm will crash due to division by zero or an empty array error.

Silhouette coefficient and elbow method: Unsupervised clustering lacks ground-truth labels, so internal metrics are needed to evaluate quality. The silhouette coefficient computes for each sample a (the average distance to other points in the same cluster) and b (the average distance to all points in the nearest different cluster), then takes (b-a)/max(a,b), with values in [-1,1]; closer to 1 means compact clusters and well-separated clusters. The elbow method plots the relationship between the number of clusters k and the within-cluster sum of squares (inertia), choosing the "elbow point" where the rate of decrease noticeably flattens as k. The two are often used together to avoid the one-sidedness of a single metric.

From Supervised to Unsupervised Learning: A Turning Point in the Learning Path

The developer's share reveals a key piece of information: they have completed all supervised learning algorithms and are now in the unsupervised learning stage, with only "2 algorithms and 2 concepts" left to finish the entire machine learning algorithm checklist. This systematic from-scratch implementation path stands in sharp contrast to the common "library-caller" route.

Supervised and unsupervised learning present different challenges in implementation difficulty. The former has clear label signals, and every step of gradient descent gets feedback from a loss function; the latter has no labels, K-Means' convergence relies on the monotonic decrease property of the objective function, and evaluating clustering quality requires indirect means such as the silhouette coefficient and elbow method. The process of writing K-Means by hand is exactly what lets you appreciate the fundamental question of unsupervised learning: how do you define a "good" result when there are no labels?

The monotonic decrease property of the objective function: K-Means' within-cluster sum of squares J = Σ∑||x - μ_i||² is monotonically non-increasing across iterations. The assignment step fixes the centers and assigns each sample to the nearest center, which only decreases or maintains J; the update step fixes the assignments and sets each center to the cluster mean, which by the optimality of the mean likewise only decreases J. Since J is bounded below by 0, the algorithm must converge to a local minimum. But this property does not guarantee global optimality, and the result depends on initialization — which is exactly why K-Means++ matters.

After Writing Algorithms by Hand: Heading Toward Deep Learning

The original poster says that after finishing the remaining algorithms and concepts they will move into deep learning, barely concealing their excitement: "Dreams have no end." From a learning-path design perspective, writing classic machine learning algorithms by hand before moving into deep learning is a solid bottom-up route. Concepts involved in K-Means — vectorized computation, objective function optimization, iterative convergence — will recur in more complex forms in deep learning: the update of cluster centers and the parameter updates of neural networks share the same "evaluate-update" loop structure in mathematical form.

For developers who want to take this path too, one pragmatic suggestion: after writing each algorithm by hand, cross-validate it against the corresponding sklearn implementation to ensure your results align on the same data with the same initialization conditions. This is both a debugging tool and a way to build confidence. Completing supervised learning plus part of unsupervised learning in 59 days is a fairly tight pace, but this kind of "slow is fast" investment will pay off when facing more complex models later.