np.argmin in kmeans now in the right axis