input: Write a function to create a k-means clustering algorithm in Python. The function should take an integer k and a dataset as input and return the cluster assignments and the final centroids.
  import numpy as np
  
  def kmeans(k, data):
      """
      Perform k-means clustering on data
      """
      # randomly set initial centroids
      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))
      
      # find the closest centroid for each sample
      while not np.allclose(centroids, prev_centroids):
          # save old centroids
          prev_centroids = centroids
      
          # get the distance between each point and each centroid
          for idx, centroid in enumerate(centroids):
              distances[:, idx] = np.linalg.norm(data - centroid, axis=1)
      
          # assign each sample to the closest centroids
          clusters = np.argmin(distances, axis=1)
      
          # update centroids 
          for idx in range(k):
              centroids[idx] = np.mean(data[clusters == idx], axis=0)
      
      return clusters, centroids
  
  # Test cases
  assert kmeans(2, np.array([[1, 2], [1, 4], [1, 0], [4, 2], [4, 4], [4, 0]]))[0].tolist() == [0, 0, 0, 1, 1, 1]
  assert kmeans(3, np.array([[1, 1], [2, 2], [3, 3], [4, 4], [5, 5], [6, 6]]))[0].tolist() == [0, 0, 1, 1, 2, 2]
output: def kmeans(k, data):
      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