KNN classifier in pure numpy

knn.py

import numpy as np 


class KNNClassifier(object):
  
  def __init__(self, data, labels, k):
    self.data = data
    self.labels = labels 
    self.k = k 
  
  def predict(sample: np.ndarray):
    """ Predicts the label of the sample based on k nearest neighbors 
        Args:
          sample: sample to classify 
    """
    dist = np.linalg.norm(self.data - sample[None, ], axis=1)
    idx = np.argpartition(dist, self.k)[:self.k]
    counts = np.bincount(self.labels[idx])
    return np.argmax(counts)