Updated KNN code