input: 
  def kmeans(k, data):
      """
      Perform k-means clustering on data.
      >>> kmeans(2, np.array([[1.0, 1.0], [2.0, 2.0], [8.0, 8.0], [9.0, 9.0]]))
      (array([0., 0., 1., 1.]), array([[1.5, 1.5], [8.5, 8.5]]))
      >>> kmeans(3, np.array([[1.0, 1.0], [2.0, 2.0], [8.0, 8.0], [9.0, 9.0], [10.0, 20.0]]))
      (array([0., 0., 1., 1., 2.]), array([[1.5, 1.5], [8.5, 8.5], [10.0, 20.0]]))
      """
      import numpy as np
      centroids = data[np.random.randint(low=0, high=len(data), size=k)]
      prev_centroids = np.zeros(centroids.shape)
      clusters = np.zeros(len(data))
      distances = np.zeros((len(data), k))
      while not np.allclose(centroids, prev_centroids):
          prev_centroids = centroids
          for idx, centroid in enumerate(centroids):
              distances[:, idx] = np.linalg.norm(data - centroid, axis=1)
          clusters = np.argmin(distances, axis=1)
          for idx in range(k):
              centroids[idx] = np.mean(data[clusters == idx], axis=0)
      return clusters, centroids
  
output:     import numpy as np
      centroids = data[np.random.randint(low=0, high=len(data), size=k)]
      prev_centroids = np.zeros(centroids.shape)
      clusters = np.zeros(len(data))
      distances = np.zeros((len(data), k))
      while not np.allclose(centroids, prev_centroids):
          prev_centroids = centroids
          for idx, centroid in enumerate(centroids):
              distances[:, idx] = np.linalg.norm(data - centroid, axis=1)
          clusters = np.argmin(distances, axis=1)
          for idx in range(k):
              centroids[idx] = np.mean(data[clusters == idx], axis=0)
      return clusters, centroids