KNN python script

import panda as pd
import numpy as np

from sklearn.neighbors import NearestNeighbors

nbrs = NearestNeighbors(n_neighbors=2, algorithm='ball_tree').fit(X_tr)

data = pd.read_csv('~/Downloads/train.csv')
test = pd.read_csv('~/Downloads/test.csv')
X_tr = data.values[:, 1:].astype(float)
y_tr = data.values[:, 0]
X_te = test.values[:, 0:].astype(float)


for i in range(0, len(X_te)):
  lst = map(lambda x:y_tr[x], nbrs.kneighbors(X_te[i].reshape(1, -1), 3, return_distance=False)[0])
  max_element = max(set(lst), key=lst.count)
  classifications.append(max_element)