input: Write a python function to assign clusters to a dataset using KMeans clustering from sklearn.
  assert assign_clusters(np.array([[2, 34], [3, 4], [7, 24], [5, 14], [8, 22], [10, 21]]), 2).tolist() == [1, 0, 1, 1, 1, 1]
  assert assign_clusters(np.array([[2, 34], [3, 4], [7, 24], [5, 14], [8, 22], [10, 21]]), 3).tolist() == [0, 2, 1, 2, 1, 1]
  assert assign_clusters(np.array([[2, 34], [3, 4], [7, 24], [5, 14], [8, 22], [10, 21], [11, 10]]), 2).tolist() == [1, 0, 1, 1, 1, 1, 0]
output: def assign_clusters(data, n_clusters):
      from sklearn.cluster import KMeans
      kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(data)
      return kmeans.labels_