rewrite KNN Algorithm (Added functionality)