Further improved kmeans code