input: Write a function to implement the k-means clustering algorithm in Python, which takes the number of clusters (k) and the dataset (data) as inputs. The function should return the cluster assignments for each data point and the final centroids.
  assert kmeans(3, np.array([[1, 2], [1, 4], [1, 0], [4, 2], [4, 4], [4, 0]])) == (np.array([0, 1, 0, 2, 1, 2]), np.array([[1, 0], [1, 4], [4, 0]]))
output: def kmeans(k, data):
      import numpy as np
      from sklearn.cluster import KMeans
      
      kmeans = KMeans(n_clusters=k, random_state=0).fit(data)
      return kmeans.labels_, kmeans.cluster_centers_