k-Nearest Neighbors From Scratch in NumPy

Given training points X_train (n × d) with integer labels y_train, predict a label for each row of X_test using Euclidean distance and the k nearest training points.

Training points at equal distance are ordered by their index in X_train, lower first.
The prediction is the label with the most votes among the k. If several labels tie, choose the one whose nearest member comes first in that order.

Coordinates are integers, so equal distances are exactly equal.

Examples

Input X_train = [[0, 0], [1, 0], [5, 5], [6, 5]], y_train = [0, 0, 1, 1], X_test = [[0.5, 0], [5.5, 5]], k = 1

Output [0, 1]

Each test point copies its single nearest neighbour.

Input X_train = [[0, 0], [2, 0], [4, 0]], y_train = [1, 2, 2], X_test = [[1, 0]], k = 3

Output [2]

k covers every point: label 2 wins two votes to one.

Input X_train = [[0, 0], [2, 0], [10, 0]], y_train = [7, 3, 3], X_test = [[1, 0]], k = 2

Output [7]

A tie at distance 1. Index order puts label 7 first, so 7 wins the 1–1 vote.

Constraints

1 ≤ k ≤ n ≤ 2000, 1 ≤ d ≤ 50, up to 500 test points
Integer coordinates in the training set

Hints

Hint 1

Build the full distance matrix with broadcasting: (n_test, 1, d) minus (1, n, d).

Hint 2

Squared distances give the same order as distances. np.argsort with kind="stable" keeps ties in index order.

Hint 3

Count votes in a dict filled in neighbour order. Python dicts remember insertion order, which is exactly the tie rule.

Where candidates go wrong

np.argsort with the default quicksort, which does not promise to keep equal distances in index order.
Breaking vote ties by the smallest label, or by Counter.most_common, whose order among equal counts is insertion order of the whole training set, not of the neighbours.
Taking a square root that changes nothing but costs time and can create float differences between equal integer distances.

The idea

Compute all squared distances at once, sort each row stably, and keep the first k indices. A stable sort leaves equal distances in index order, which is the first tie rule. Count votes in a dictionary filled in neighbour order: among labels with the top count, the first key in the dictionary is the one whose nearest member came first, which is the second tie rule.

With integer coordinates every squared distance is an exact integer in float64, so ties really are ties.

Time O(n_test · n(d + log n)) · Space O(n_test · n · d) with broadcasting

Reference solution

import numpy as np

def knn_predict(X_train, y_train, X_test, k):
    A = np.asarray(X_train, dtype=float)
    B = np.asarray(X_test, dtype=float)
    y = np.asarray(y_train)
    D = ((B[:, None, :] - A[None, :, :]) ** 2).sum(axis=2)   # squared distances, (n_test, n)
    out = []
    for row in D:
        near = np.argsort(row, kind='stable')[:k]           # equal distances stay in index order
        votes = {}
        for j in near:                                      # filled nearest first
            votes[y[j]] = votes.get(y[j], 0) + 1
        top = max(votes.values())
        out.append(int(next(label for label in votes if votes[label] == top)))
    return out

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

kNN is the simplest model there is, so the interview is about engineering: a vectorised distance matrix, the choice of squared distance, and clear rules for ties. Candidates who never think about ties fail the hidden tests.

The theory behind it is in the lesson k-means, kNN, PCA and decision trees from scratch.

Keep going

Easy12 min

k-nearest neighbours classifier

Examples

Constraints

Hints

Where candidates go wrong

The idea

Reference solution

Follow-ups an interviewer asks

Why interviewers ask it

Keep going

k-nearest neighbours classifier

k-means clustering (Lloyd's algorithm)

Principal component analysis

Best decision-tree split by Gini impurity