Function change to fit new version of Kmeans