k-Nearest Neighbours From Scratch

Question

Create a function that classifies a new observation using the k-nearest neighbours algorithm. You are not allowed to use any machine learning libraries (so no scikit-learn).
def knn_predict(X_train: List[List[float]], y_train: List[int], x_new: List[float], k: int) -> int:
    pass
Input
X_train is a list of training points (each a list of floats), y_train contains the integer class label of each training point, x_new is the point to classify, and k is the number of neighbours to use (1 ≤ k ≤ len(X_train)).

Output
Return the predicted class label as an integer.

Bonus
Handle the case where two classes receive the same number of votes.

Example Tests

 	With X_train = [[1.0, 1.0], [1.5, 2.0], [3.0, 4.0], [5.0, 7.0], [3.5, 5.0], [4.5, 5.0]] and y_train = [0, 0, 1, 1, 1, 1]:
 	knn_predict(X_train, y_train, [1.2, 1.5], 3) should return 0.
 	knn_predict(X_train, y_train, [4.0, 5.0], 3) should return 1.

Accepted Answer

k-Nearest Neighbours classifies a new point by looking at the $k$ training points that lie closest to it and letting them vote. The algorithm has three steps:

 	Distancecompute the Euclidean distance between $x_{new}$ and every training point: \begin{equation} d(a, b) = \sqrt{\sum_{j=1}^{d} (a_j - b_j)^2} \end{equation}
 	Selectsort the training points by this distance and keep the $k$ smallest.
 	Votecount the class labels among these $k$ neighbours and return the most common one.

For the bonus (a tie in the vote), a clean rule is: among the tied classes, pick the class of the nearest neighbour. Since we already sorted by distance, we can simply walk through the neighbours in order and return the first label that belongs to a tied class.
import math
from collections import Counter
from typing import List

def knn_predict(X_train: List[List[float]], y_train: List[int], x_new: List[float], k: int) -> int:
    # Step 1: compute the Euclidean distance to every training point
    distances = []
    for features, label in zip(X_train, y_train):
        dist = math.sqrt(sum((a - b) ** 2 for a, b in zip(features, x_new)))
        distances.append((dist, label))

    # Step 2: sort by distance and keep the k closest points
    distances.sort(key=lambda pair: pair[0])
    nearest = distances[:k]

    # Step 3: majority vote among the k neighbours
    votes = Counter(label for _, label in nearest)
    max_votes = max(votes.values())
    tied = [label for label, count in votes.items() if count == max_votes]
    if len(tied) == 1:
        return tied[0]

    # Bonus: tie-break by the class of the nearest neighbour among the tied classes
    for _, label in nearest:
        if label in tied:
            return label

# Test the function with the given test cases
X_train = [[1.0, 1.0], [1.5, 2.0], [3.0, 4.0], [5.0, 7.0], [3.5, 5.0], [4.5, 5.0]]
y_train = [0, 0, 1, 1, 1, 1]
print(knn_predict(X_train, y_train, [1.2, 1.5], 3))  # 0
print(knn_predict(X_train, y_train, [4.0, 5.0], 3))  # 1

A few remarks that score points in an interview:

 	Complexityfor $n$ training points in $d$ dimensions, one prediction costs $O(nd)$ for the distances plus $O(n \log n)$ for the sort. (Using a partial selection instead of a full sort brings the second part down to $O(n \log k)$.) There is no training cost at all — kNN is a lazy learner.
 	Feature scalingbecause the prediction is built entirely on distances, features on large scales dominate the vote. Standardise features before using kNN.
 	Choice of ksmall $k$ gives high variance, large $k$ gives high bias; choose $k$ with cross-validation, and prefer an odd $k$ for binary problems to reduce ties.

Title	Category	Subcategory	Difficulty	Status
How a Decision Tree Chooses a Split	Machine Learning	Classic Algorithms	Easy	Unlock
kNN vs k-Means	Machine Learning	Classic Algorithms	Easy	Unlock

The Interview Guide for Quants and Traders

k-Nearest Neighbours From Scratch

What do you think of this page?