method predict

def euclidian_distance(point_a, point_b):
    return (sum([(point_a[i] - point_b[i])**2 for i in range(len(point_a))]))**0.5

class KMeans:
    def __init__(self, k):
        self.k = k

    def find_cluster_centroid(self, cluster_points):
    # Write your code here
        centroid = []
        for point in cluster_points:
            for dim, value in enumerate(point):
                if len(centroid) - 1 < dim:
                    centroid.append(0)
                centroid[dim] += value
        for i in range(len(centroid)):
            centroid[i] /= len(cluster_points)
        return centroid

    def fit(self, X_train):
        # Write you code here
        self.centroids = [X_train[i] for i in range(2, self.k*2 + 1, 2)]
        while True:
            # Assing each point to centroid
            points_clusters = []
            for point in X_train:
                distances = []
                for centroid in self.centroids:
                    dist_to_point = euclidian_distance(centroid, point)
                    distances.append(dist_to_point)
                closest_centroid_index = distances.index(min(distances))
                points_clusters.append((point, closest_centroid_index))

            # Update centroids
            clusters = {}
            for point, cluster in points_clusters:
                if cluster not in clusters:
                    clusters[cluster] = []
                clusters[cluster].append(point)
            new_centroids = [None]*self.k
            for cluster, points in clusters.items():
                new_centroids[cluster] = self.find_cluster_centroid(points)

            # Check if converged
            total_dist = 0
            for i in range(len(new_centroids)):
                total_dist += euclidian_distance(new_centroids[i], self.centroids[i])
            if total_dist < 0.01:
                break
            self.centroids = new_centroids


    def predict(self, X_test):
        pass