How do you write k-means and k-NN from scratch in NumPy, step by step?

In this article

1. Before you start
2. Why interviewers ask this
3. K-means and k-NN from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page belongs to the machine learning (ML) coding round guide, where from-scratch tasks like this one are collected.

Why interviewers ask this

A from-scratch k-means or k-NN task is reported by candidates in coding rounds at large tech companies, often with the follow-up of how to vectorize the code. The task is small enough to finish in one sitting, yet it exposes whether you can turn an idea into array operations.

A weak answer writes three nested loops and cannot say what the shapes are. A strong answer builds the distance matrix with broadcasting, states the stopping rule, and knows that k-means can stop at a poor answer. A Staff-level answer adds restarts, feature scaling, memory cost of the broadcast, and the difference between the two algorithms: one is unsupervised, the other supervised.

The two are often confused because both use the letter k. They mean different things, and a clear first sentence about that difference earns trust before any code is written.

K-means and k-NN from first principles

k-means++

k-means++ is a seeding scheme that chooses the initial centroids by sampling from a probability distribution based on each point's contribution to inertia, which the scikit-learn KMeans page says speeds up convergence.

Broadcasting

Broadcasting is the NumPy rule that lets arrays of different shapes be combined by stretching each axis of size one to match the other array, so one expression can compute all pairwise differences without a Python loop.

In plain words

Picture sorting a pile of coins into two heaps. Drop two markers anywhere. Push each coin to the nearer marker. Move each marker to the middle of its heap. Repeat until nothing changes. That is k-means. The analogy breaks in one place: real data has many features, not two table coordinates, and the markers can settle in a poor spot.

For k-NN, imagine asking the few neighbors nearest to your house how they voted, and following the majority. In one precise sentence each: k-means needs the number of clusters given up front and minimizes inertia, and k-NN is instance-based learning, which stores the training data and predicts by a majority vote of the nearest neighbors.

Step 1: Pick k starting centroids

K-means needs k starting centroids. The simplest start copies k randomly chosen rows of the data, which is what scikit-learn's undefined initialization does. The better-known start is k-means++, defined above, which spreads the starting centroids by probability instead of copying random rows.

Input: the data array X of shape (n, d) and the integer k. Output: a centroid array C of shape (k, d). In the examples below we use six points, [[1,2],[1,4],[1,0],[10,2],[10,4],[10,0]], the example in the scikit-learn k_means reference, and start the centroids on the first two rows. Starting position matters a great deal, as the trade-offs section shows.¹

Step 2: One broadcast builds every distance at once

The distance matrix D of shape (n, k) holds the squared distance from every point to every centroid. Insert a size-one axis so that X[:, None, :] has shape (n, 1, d) and C[None, :, :] has shape (1, k, d). Subtracting them broadcasts to shape (n, k, d); squaring and summing over the last axis leaves (n, k).

For our data the shapes are (6, 1, 2) and (1, 2, 2), which broadcast to (6, 2, 2) and collapse to (6, 2).¹ The first row of D is (0, 4): the point (1, 2) is on top of the first centroid and 4 away, squared, from the second.¹ Skipping the square root is safe because the nearest centroid by squared distance is the nearest by distance. This is the single line interviewers mean by "vectorize it", and the next two steps reuse it.

import numpy as np

def sq_dists(A, B):
    # A: (n, d), B: (m, d)  ->  (n, m) squared Euclidean distances
    return ((A[:, None, :] - B[None, :, :]) ** 2).sum(axis=2)

Step 3: Assign every point, then re-average each cluster

Assignment is labels = D.argmin(axis=1): for each row, the index of the nearest centroid. The update step replaces each centroid with the mean of the points carrying its label. Input: D and X. Output: new labels and new centroids.

An empty cluster has no points to average. One common choice, used below, is to keep its old centroid. After one pass on our six points, the centroids move from (1, 2) and (1, 4) to (5.5, 1) and (5.5, 4), and the inertia, computed with the starting centroids, is 251.¹ Each point is now closer to a better-placed centroid, and the next pass checks whether anything moves.

Step 4: Stop when the centroids stop moving

Repeat steps 2 and 3 until the centroids no longer change, or until a cap on the iteration count is reached. scikit-learn's defaults are a cap of 300 iterations and a tolerance of 0.0001. Its documentation calls this loop the classical "lloyd" algorithm, and it always ends because each pass can only lower or keep the inertia.

Return labels and the inertia, which map to scikit-learn's labels_ and inertia_ attributes. The full function is short:

def kmeans(X, init_idx, max_iter=100):
    C = X[init_idx].astype(float)
    for _ in range(max_iter):
        D = sq_dists(X, C)
        labels = D.argmin(axis=1)
        newC = C.copy()
        for j in range(len(C)):          # loop over k, not over n
            if (labels == j).any():      # keep the old centroid if empty
                newC[j] = X[labels == j].mean(axis=0)
        if np.allclose(newC, C):
            break
        C = newC                         # if max_iter ends the loop, C is one update ahead of labels
    inertia = D[np.arange(len(X)), labels].sum()
    return C, labels, inertia

The only Python loop runs over the k clusters, never over the n points. When the loop stops because the centroids stopped moving, the inertia is measured with the centroids used for the last assignment, so it matches the labels returned; if max_iter runs out first, the returned centroids are one update ahead of the labels and the inertia.

Step 5: For k-NN, sort the distances once and vote

K-NN reuses sq_dists, now between query points Q and the stored training set. Take the k smallest distances per row with np.argsort(D, axis=1)[:, :k], look up their labels, and return the most common one per row with np.bincount(...).argmax(). To find only the k nearest, np.argpartition(D, k, axis=1)[:, :k] does it in linear time per row where a full sort costs n log n, at the price of returning them unsorted.

def knn_predict(X_train, y_train, Q, k):
    D = sq_dists(Q, X_train)                  # (queries, train)
    idx = np.argsort(D, axis=1)[:, :k]        # k nearest per query
    votes = y_train[idx]                      # (queries, k) labels
    return np.array([np.bincount(v).argmax() for v in votes])

K-NN sorts the distances once and takes a majority vote of the k nearest labels. K-NN stores the training data and pays its cost at prediction time. Nothing is trained: the function only stores and compares, the instance-based behavior the documentation describes. In the vote, bincount(...).argmax() breaks a tie toward the lowest label, a convention you should say aloud. The input is the query array and the stored data; the output is one predicted label per query. You now have both algorithms, built from one shared distance function.

1 · Start2 · Distances3 · Assign4 · Stop5 · k-NN voteReplay

Step-by-step description

1. Pick k starting centroids.
2. One broadcast gives every distance.
3. Assign points, then re-average.
4. Stop when centroids stop moving.
5. Sort distances once and vote.

• Broadcasting builds the distance matrix of shape (n, k) in one expression.
• K-means assigns each point and re-averages the centroids until they stop moving.
• K-NN sorts the distances once and takes a majority vote of the k nearest labels.

Worked example

We run the code on two small data sets. Every number below is computed by the script, not typed by hand.¹

First, k-means on the six points with k = 2 and starting centroids at the first two rows, (1, 2) and (1, 4). Pass one gives the labels (0, 1, 0, 0, 1, 0). The centroids become (5.5, 1) and (5.5, 4). Pass two assigns the same labels, so the loop stops with an inertia of 125.5.¹

Second, the same data starting at (1, 2) and (10, 2). The labels are (0, 0, 0, 1, 1, 1), the centroids stay put, and the inertia is 16.0, the value the scikit-learn k_means reference shows for this data. Check by hand: each cluster contributes 0 + 4 + 4 = 8, and 8 + 8 = 16.

For k-NN, take training points (1,1), (2,1), (1,2) with label 0 and (5,5), (6,5), (5,6) with label 1 (illustrative), and the query (3, 3). The squared distances are 8, 5, 5, 8, 13, 13. The third-nearest place is a tie between (1, 1), label 0, and (5, 5), label 1, so the two nearest, (2, 1) and (1, 2), both label 0, are joined by one of the tied points; either choice gives at least two votes of three for label 0, and the prediction is 0.¹ np.argsort with its default sort is not stable, so pass kind="stable" for reproducible tie-breaks.

Text description of the diagram

• The best inertia is 137.5 at k=1 and 16 at k=2.
• Inertia reaches 0 when each of the six points is its own cluster.

The bars show the best inertia over every choice of starting rows. It reaches zero when each point is its own cluster, which matches the documented fact that inertia is not normalized: lower is better and zero is optimal. A basic evaluation script therefore plots inertia against k and looks for where the drop flattens, rather than picking the lowest value.

Trade-offs and edge cases

K-means can stop at a poor answer, as the 125.5 run shows. Two documented fixes exist. Restart from several starts and keep the best: scikit-learn returns the best of undefined runs by inertia, where n_init='auto' means 10 runs for random starts and 1 for k-means++. Or seed smartly with k-means++.

K-means also assumes convex, isotropic clusters and responds poorly to elongated or irregular shapes. In very high dimensions Euclidean distances inflate, and the same guide suggests reducing dimensions first, for example with principal component analysis (PCA).

K-NN has its own traps. Feature scaling: with age in years and income in dollars (illustrative), a query's nearest neighbor by raw distance is a point 30.48 away that differs mainly in age; after standardizing each feature, a different point is nearest, at a distance of 0.204, and the predicted label flips.¹ Cost: the brute-force approach becomes infeasible as the number of samples grows, and tree-based search is the alternative. The best k is highly data-dependent, so tune it on held-out data. Voting can also weight neighbors by distance, as KNeighborsClassifier's undefined option allows.

Memory is the vectorization cost. The broadcast tensor holds n × k × d numbers: for 10,000 points, 10 clusters and 50 features (illustrative) that is 5,000,000 numbers, about 40 MB in 64-bit floats. For k-NN with 1,000 queries against 100,000 stored points (illustrative), the distance matrix has 100,000,000 entries, about 0.8 GB.

The page's sq_dists first builds the broadcast difference tensor of shape (queries, train, d), which is 1,000 × 100,000 × 50 = 5,000,000,000 numbers, about 40 GB in 64-bit floats, so the 0.8 GB figure holds only for the finished matrix. The standard fix expands the square: squared distance = ‖q‖² − 2 q·x + ‖x‖², which forms only the (queries, train) matrix through one matrix product (clip tiny negative values caused by rounding to zero). Process queries in chunks either way.¹

Text description of the diagram

• K-means is unsupervised and k-NN is supervised.
• K-NN stores the training data and pays its cost at prediction time.

For other from-scratch drills, see gradient descent from scratch and the numerically stable softmax.

How to answer it in an interview

The spoken structure: name the two steps, state the shapes, give the stopping rule, then name the local-minimum risk and its fix.

Common mistakes

• "K-means and k-NN both learn from labels." (wrong) K-means has no labels; k-NN needs them. (right) "K-means is unsupervised and k-NN is supervised; the k means clusters in one and voters in the other."
• "I loop over every point and every centroid." (wrong) That is the version that was asked to be vectorized. (right) "One broadcast gives an n-by-k matrix; only the loop over k clusters remains."
• "Lower inertia means the better k." (wrong) Inertia falls as k grows, reaching zero at one point per cluster. (right) "I compare inertia across k and look where the gain flattens."
• "K-NN needs no preparation, so scaling is irrelevant." (wrong) A large-unit feature dominates the distance. (right) "I standardize features before computing distances."
• "K-means finds the best clustering." (wrong) The poor-start run stopped at 125.5, not 16.0. (right) "It finds a local answer; I restart and keep the lowest inertia."

Follow-up questions

How would you vectorize this k-means implementation?

Replace the nested loops over points and centroids with the broadcast X[:, None, :] - C[None, :, :], square and sum over the last axis, then call argmin(axis=1). The update step keeps one loop over k, or can use a one-hot matrix product.

How do you write a basic evaluation script for a clustering implementation?

Return the inertia and plot it for a range of k. Check your function on a data set with a known answer, such as the six points above, where the inertia must be 16.0 for the two natural groups. Compare against scikit-learn's KMeans on the same data, comparing partitions rather than label ids, because label numbering depends on the seed.

How do you implement a k-nearest-neighbors classifier using only NumPy?

Build the query-by-train squared-distance matrix, take argsort along axis one, keep the first k columns, look up their labels, and return the most frequent label per row with bincount.

Why does k-means give different answers on different runs?

The starting centroids are random, and the loop stops at the nearest settling point, so different starts can end at different inertia values (125.5 against 16.0 above).¹ Fix the random seed for reproducibility, and use restarts to find a lower inertia.

Does k-NN need feature scaling?

Yes in practice, because the prediction rests on distance, and a feature with large units dominates it. The scaling example above flips the predicted label after standardizing.¹

What changes for a much larger data set?

K-NN prediction cost grows with the stored data, so tree-based or approximate search replaces brute force. Before that, replace the broadcast tensor with the expanded square, ‖q‖² − 2 q·x + ‖x‖², so memory scales with the (queries, train) matrix rather than with the feature count too, and process queries in chunks. For related from-scratch topics, see scaled dot-product attention, multi-head attention, 2D convolution and building a generative pre-trained transformer (GPT).

Sources

1. 2.3. Clustering: K-means — scikit-learn documentation (accessed 2026-10-26), 26 Oct 2026
2. sklearn.cluster.KMeans — scikit-learn documentation (accessed 2026-10-26), 26 Oct 2026
3. sklearn.cluster.k_means — scikit-learn documentation (accessed 2026-10-26), 26 Oct 2026
4. 1.6. Nearest Neighbors — scikit-learn documentation (accessed 2026-10-26), 26 Oct 2026
5. sklearn.neighbors.KNeighborsClassifier — scikit-learn documentation (accessed 2026-10-26), 26 Oct 2026