Implement k-NN Efficiently

Math Foundations

Given training set $(X_{\text{train}}, y_{\text{train}})$ and query set $X_{\text{test}}$, k-Nearest Neighbors (k-NN):

1. Computes Euclidean Distance from each query sample $x_q$ to all training samples $x_i$:

$$d(x_q, x_i) = \sqrt{\sum_{j=1}^D (x_{q,j} - x_{i,j})^2}$$

2. Identifies the $K$ smallest distance neighbors.
3. Assigns the majority vote class label (Classification) or mean target value (Regression).

Matrix Expansion Speed Trick

Calculating pairwise distances using nested Python loops is extremely slow ($\mathcal{O}(N_{\text{test}} \times N_{\text{train}} \times D)$).

Using the algebraic expansion of squared Euclidean distance:

$$| Q - T |^2 = | Q |^2 + | T |^2 - 2 (Q T^T)$$

• $| Q |^2$: Sum of squares along test rows ([N_test, 1]).
• $| T |^2$: Sum of squares along train rows ([1, N_train]).
• $Q T^T$: High speed matrix multiplication (np.dot(Q, T.T)).

This computes all $N_{\text{test}} \times N_{\text{train}}$ pairwise distances using optimized C / BLAS matrix math!

Efficient NumPy Implementation

import numpy as np


class KNNClassifierEfficient:

  def __init__(self, k=5):
    self.k = k
    self.X_train = None
    self.y_train = None

  def fit(self, X, y):
    # Lazy learning: store training data
    self.X_train = X
    self.y_train = y

  def _compute_distances_fast(self, X_test):
    # Fast pairwise squared Euclidean distance using matrix expansion
    # D^2 = sum(Q^2) + sum(T^2) - 2 * (Q @ T.T)
    test_sum = np.sum(np.square(X_test), axis=1, keepdims=True)  # [N_test, 1]
    train_sum = np.sum(
        np.square(self.X_train), axis=1, keepdims=True
    ).T  # [1, N_train]
    dot_product = np.dot(X_test, self.X_train.T)  # [N_test, N_train]

    dists_sq = test_sum + train_sum - 2 * dot_product
    # Clip small negative values due to floating point imprecision
    return np.sqrt(np.maximum(dists_sq, 0.0))

  def predict(self, X_test):
    dists = self._compute_distances_fast(X_test)
    num_test = X_test.shape[0]
    y_pred = np.zeros(num_test, dtype=self.y_train.dtype)

    for i in range(num_test):
      # np.argpartition finds K smallest indices in O(N) time!
      knn_indices = np.argpartition(dists[i], self.k)[: self.k]
      knn_labels = self.y_train[knn_indices]

      # Majority vote
      counts = np.bincount(knn_labels)
      y_pred[i] = np.argmax(counts)

    return y_pred

Say this out loud

Implementing k-NN efficiently utilizes matrix expansion to compute all pairwise Euclidean distances using high speed matrix multiplication. Using np.argpartition identifies top K nearest neighbor indices in linear time without doing a full array sort. Feature scaling is mandatory before running k-NN.

Followups to expect

1. What is the computational complexity of k-NN inference? Brute force k-NN inference takes $\mathcal{O}(N_{\text{test}} \times N_{\text{train}} \times D)$ time. Using tree indexes (KD-Tree or Ball Tree) reduces query search time to $\mathcal{O}(D \log N_{\text{train}})$.
2. Why does k-NN struggle in high dimensional spaces? The Curse of Dimensionality causes distances between all pairs of points in high dimensional space to converge to near identical values, ruining neighbor discrimination.