ML Coding From Scratch

How would you write a KNN search using base Python?

By Interview Clozo · Published 27 October 2026 · Last verified 26 October 2026 · 12 min read

How would you write a k-nearest neighbours (KNN) search using base Python? Write a loop that measures the distance from a query to every stored point and returns the k closest; a KNN classifier then takes a majority vote of their labels. Only the standard library is needed. One slip silently changes the neighbours: features on different scales, shown in the second step.

In this article

Before you start
Why interviewers ask this
KNN search from first principles
The Staff-level answer
The rubric
Three variants of this question
Common mistakes
Follow-up questions
Sources

Before you start

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

Why interviewers ask this

A KNN search is a small problem that shows whether a candidate can turn an idea into correct code without a library to lean on. The loop is short, so the interviewer's attention goes to the decisions around it: which distance, whether to sort everything, what happens on a tie, and what happens when one feature is measured in thousands and another in single digits.

A weak answer writes a nested loop, sorts all the distances and returns the most common label, with no remark on cost or scaling. A strong answer selects the k nearest without a full sort, handles ties on purpose and says what its cost is. A Staff-level answer also says when brute force stops being acceptable and which index replaces it. In a design or coding round, a prompt like this one tests that range of judgement as much as the syntax. Beam search uses the same keep-the-best-k idea on partial sequences, so the selection step here transfers.

KNN search from first principles

Brute-force nearest-neighbour search: Brute-force nearest-neighbour search is a method that measures the distance from a query to every stored point and then picks the smallest, so it needs no index and no training step but touches every point on every query.

In plain words

Imagine asking the five houses nearest to yours what colour they are painted. Finding those houses is a KNN search. Guessing your own colour from the most common answer is the classifier built on top. The analogy breaks in one place: houses sit on a flat map. Data points can have many features, and "near" then depends on how each feature is measured.

In one precise sentence: given a query, a KNN search finds the k stored points with the smallest distance to it, and a classifier returns the majority label among them. Brute force builds nothing before the first query; the tree variants build an index first, as the scikit-learn guide describes.

The worked example uses six labelled points in two dimensions (illustrative): A (1, 1), B (2, 1) and C (1, 3) are red, D (5, 5), E (6, 4) and F (7, 6) are blue. The query is (3, 2).¹

Step 1: Store the labelled points

Step 1 stores the training data as a list of (features, label) pairs. Nothing is computed yet: brute-force KNN has no training phase, and "fitting" it means keeping the list. Input: the labelled points. Output: the same list, kept in memory.

Write a point as a tuple, for example ((1, 1), "red"). A tuple suits a point because it cannot be changed by accident while you loop over it. Keep every point the same length, because a distance is only defined between points of the same dimension, as the math.dist documentation states.

You now know what the search runs over: a plain list. Its size is called N, and the number of features per point is called D. Both appear in the cost at the end of the page.

Step 2: Measure the distance to every stored point

Step 2 computes the distance from the query to each stored point, which costs work proportional to N times D per query: 200,000 coordinate differences for N of 100,000 and D of 2 (illustrative).¹ In base Python, math.dist(p, q) returns the Euclidean distance and has existed since Python 3.8; sum((a - b) ** 2 for a, b in zip(p, q, strict=True)) needs no import.

Squared distance gives the same ordering as the true distance, because the square root never swaps two values, so skip the root. For the example, the true distances are 2.2361 (A), 1.4142 (B) and 2.2361 (C). D and E tie at 3.6056 and F is farthest at 5.6569.¹

Now the slip from the lead. Take a query (30, 50,000) for an age and an income, with three stored points: red (30, 58,000), blue (60, 50,500) and blue (45, 70,000) (illustrative). Raw squared distances to the first two are 64,000,000 for red and 250,900 for blue, so blue looks nearest only because income is measured in thousands. Fit min-max scaling on the stored points only (age 30 to 60, income 50,500 to 70,000), then apply the same scaling to the query: red is at about 0.4103 and blue at 1.0003, so red wins.¹ Scale features before measuring.

Step 3: Keep only the k smallest distances

Step 3 selects the k nearest points without sorting all N distances. The simple route is sorted(...)[:k]; the alternative is heapq.nsmallest(k, train, key=...), which the Python 3.1 documentation describes as equivalent to sorted(iterable, key=key)[:n] while returning only the n smallest. The costs differ. Measuring all distances is O(N·D), a full sort adds O(N log N), and nsmallest adds O(N log k) time with O(k) extra memory. The same Python 3.1 heapq page adds a caveat: these functions "perform best for smaller values of n", sorted() is more efficient for larger values, and min() is better when only one item is needed.

A heap is a binary tree stored in a list in which every parent is smaller than or equal to its children, so the smallest element is always at heap[0] (Python heapq documentation). To keep the k smallest, hold a heap of size k holding the k best so far with the worst at the root; Python's heap is a min-heap, so store negated distances (the max-heap helpers exist only from Python 3.14). A new point enters only if it beats the root.

For 100,000 random distances and k of 5, the heap changed only 50 times, against about 54 expected for randomly ordered data, and the result matched the sorted answer.¹ Most points are rejected after one comparison.

On the example, the order is B, A, C, then D and E tied, then F. The three nearest, B, A and C, are all red.¹

A KNN search is four moves: store, measure, select the k nearest, vote.CC BY 4.0 — source: How would you write a KNN search using base Python?

Step-by-step description

Keep the labelled points in a list.
Compute the distance to every point.
Keep the k smallest distances.
The most common label wins.

The labelled points are stored in a list.
The distance from the query to every point is measured.
The k smallest distances are kept.
A majority vote of their labels gives the answer.

Step 4: Take a majority vote of the k labels

Step 4 turns the k nearest labels into one answer. A majority vote is the rule that the label carried by the most neighbours wins; in base Python, collections.Counter counts the labels and the largest count decides.

With k of 3 the example votes red three to zero. With k of 4 the fourth neighbour is blue, one of two tied points D and E at the same distance, and the vote is three to one for red. With k of 5 it is three to two.¹ Tied distances at the boundary of k mean "the k nearest" is not unique, so state a rule.

An even k can also split two classes evenly. Pick an odd k for two classes, or break the tie by the closest neighbour, which the code below does. This is the point where a candidate shows the problem was understood.

The whole function in base Python

The code below has two functions. knn_search returns the k nearest (features, label) pairs, and knn_classify votes on them. It uses squared distance, heapq.nsmallest and a tie-break that prefers the label of the closest neighbour, because nsmallest returns its items in ascending order.

import heapq
from collections import Counter

def sq_dist(a, b):
    # squared Euclidean distance; same order as the true distance
    return sum((x - y) ** 2 for x, y in zip(a, b, strict=True))

def knn_search(train, query, k):
    # train: list of (features_tuple, label); features already scaled
    return heapq.nsmallest(k, train, key=lambda item: sq_dist(item[0], query))

def knn_classify(train, query, k):
    nearest = knn_search(train, query, k)
    votes = Counter(label for _, label in nearest)
    top = max(votes.values())
    for _, label in nearest:      # nearest first, so ties go to the closest neighbour
        if votes[label] == top:
            return label

strict=True makes zip raise an error if the query and a stored point have different lengths, instead of silently ignoring the extra coordinates.

Calling knn_classify on the six example points with the query (3, 2) and k of 3 returns "red", matching the hand result above.¹

The three shortest bars belong to red points, so the k=3 vote is red.CC BY 4.0 — source: How would you write a KNN search using base Python?

Text description of the diagram

B is nearest at 1.4142.
A and C tie at 2.2361.
D and E tie at 3.6056 and F is farthest at 5.6569.

Knowing it is half the job. Saying it is the other half.

Interview Clozo is an AI interview copilot for real interview rounds: it hears the question and puts a strong answer on your screen, and you say it in your own words. See the features or how it works.