input: Write a python function to perform KMeans clustering on a given dataset and return the cluster labels.
  assert k_means_clustering([[1, 2], [1, 4], [1, 0], [4, 2], [4, 4], [4, 0]], 2).tolist() == [0, 0, 0, 1, 1, 1]
  assert k_means_clustering([[1, 2], [1, 4], [1, 0], [4, 2], [4, 4], [4, 0]], 3).tolist() == [0, 1, 0, 2, 1, 2]
  assert k_means_clustering([[1, 2], [1, 4], [1, 0], [4, 2], [4, 4], [4, 0]], 1).tolist() == [0, 0, 0, 0, 0, 0]
output: def k_means_clustering(data, k):
      from sklearn.cluster import KMeans
      kmeans = KMeans(n_clusters=k, random_state=0).fit(data)
      return kmeans.labels_