numba_knn.py

from numba import jit, njit
@jit('float64[:](float64[:], float64[:, :])')
def comp_inner_numba(i, x):
    res = np.zeros(x.shape[0], dtype=np.float64)
    for j in range(x.shape[0]):
        res[j] = np.sqrt(np.sum((i-x[j])**2))
    return res

def nearest_ngbr_numba(x):
    dist = {}
    for idx,i in enumerate(x):
        lst = comp_inner_numba(i,x)
        s = np.argsort(lst)
        sorted_array = np.array(x)[s][1:]
        dist[idx] = s[1:]
    return dist