Vectorized K-nearest neighbors (no for-loop)

knn-vectorized.py

import numpy as np
from scipy.stats import mode


def euclidean_distance(x1, x2):
    # x1 array of shape (N, D)
    # x2 array of shape (N', D)
    # output pairwise distances, array of shape (N, N')
    return np.linalg.norm(x1[:, np.newaxis] - x2, axis=1)


class KNN:
    def __init__(self, k=3):
        self.k = k

    def fit(self, x, y):
        self.x_train = x
        self.y_train = y

    def predict(self, x):
        distances = euclidean_distance(x, self.x_train)
        topk_idxs = np.argsort(distances, axis=1)[:, :self.k]
        topk_labels = self.y_train[topk_idxs]
        return mode(topk_labels, axis=1)
      

if __name__ == "__main__":
    n1 = 24
    n2 = 7
    d = 3
    k = 3
  
    x = np.random.rand(n1, d)
    y = np.random.randint(0, 2, n1)
  
    knn = KNN(k)
    knn.fit(x, y)
  
    x = np.random.rand(n2, d)
    print(knn.predict(x))