input: def kmeans(k, data): 
      """
      Perform k-means clustering on data
      >>> kmeans(3, np.array([[1, 2], [3, 4], [2, 1], [4, 3], [5, 5], [1, 4]]))
      (array([0, 1, 0, 1, 2, 0]), array([[2.        , 2.        ],
                                        [3.5       , 3.5       ],
                                        [5.        , 5.        ]]))
      """
output:     import numpy as np
      
      def euclidean_distance(a, b):
          return np.linalg.norm(a - b)
      
      def initialize_centroids(data, k):
          return data[np.random.choice(data.shape[0], k, replace=False)]
      
      def assign_clusters(data, centroids):
          clusters = np.zeros(data.shape[0], dtype=int)
          for i, point in enumerate(data):
              distances = [euclidean_distance(point, centroid) for centroid in centroids]
              clusters[i] = np.argmin(distances)
          return clusters
      
      def update_centroids(data, clusters, k):
          new_centroids = np.zeros((k, data.shape[1]))
          for i in range(k):
              cluster_points = data[clusters == i]
              if cluster_points.size > 0:
                  new_centroids[i] = np.mean(cluster_points, axis=0)
          return new_centroids
      
      centroids = initialize_centroids(data, k)
      clusters = assign_clusters(data, centroids)
      
      while True:
          new_centroids = update_centroids(data, clusters, k)
          new_clusters = assign_clusters(data, new_centroids)
          if np.array_equal(clusters, new_clusters):
              break
          clusters = new_clusters
          centroids = new_centroids
      
      return clusters, centroids