Fix bugs in k-NN implementation