Optimize K-Mean for large number of clusters I am writing a Python code for KMeans clustering.
The aim of this post is to find out how I can make my below mentioned code optimal when the number of clusters is very large. I am dealing with data with tens of millions of samples and scores of features. The number of clusters that I am talking about is of the order of 500000. I have checked the correctness of the code by comparing that with results from scikit-learn's implementation of KMeans clustering.
I am aware that I can use sckit learn for KMeans clustering. But, I am writing my own custom Python code for a few purposes which include but are not limited to the below reasons.

To be able to change distance functions for my specific work assignments : from Euclidean distance to Huang distance.
To use this custom code for various other tasks like anomaly detection.

Below is my code:
class KMeans:
    def __init__(self, arr, k):
        self.arr = arr
        self.k = k

    # Main driver function
    def loop_over_centers(self, max_iter):

        # Choose random centroids.
        random_centers = np.random.choice(len(self.arr), self.k, replace=False)

        # Iterate till a certain max number of iterations.
        for iteration in range(0, max_iter, 1):
            cluster_dict = defaultdict(list)

            if iteration == 0:
                chosen_centers = random_centers
                

            # Assign samples to cluster centers based on minimum distance
            for i in range(len(self.arr)):
                distance_list = []
                for center in chosen_centers:
                    distance_list += [self.distance(self.arr[center], self.arr[i])]

                cluster_dict[chosen_centers[np.argmin(distance_list)]] += [i]

            chosen_centers = self.update_center(cluster_dict)

        return cluster_dict


    # Distance between two arrays
    def distance(self, arr1, arr2):
        return np.sqrt(np.sum(np.subtract(arr1, arr2) ** 2))


    # Gives new centroids with each new iteration
    def update_center(self, cluster_dict):
        cluster_center = []

        for cluster in list(cluster_dict.keys()):
            cluster_list = []

            # Find samples corresponding to each cluster.
            for sample in cluster_dict[cluster]:
                cluster_list += [self.arr[sample]]
                

            # Mean of samples in a given cluster.
            cluster_mean = np.mean(np.array(cluster_list), axis=0).tolist()

            distance_from_mean = []

            # Find which sample is closest to mean of cluster.
            for sample in cluster_dict[cluster]:
                distance_from_mean += [self.distance(cluster_mean, self.arr[sample])]


            # Build an updated list of cluster centers.
            cluster_center += [cluster_dict[cluster][np.argmin(distance_from_mean)]]

        return cluster_center


arr = [[12, 3, 5], [1, 6, 7], [8, 92, 1], [6, 98, 2], [5, 90, 3], [29, 5, 6], [11,4,4], [30, 6, 5],]

testKMeans = KMeans(arr, 3)
testKMeans.loop_over_centers(1000)

In a real world case, arr will have tens of millions of samples and I would like to have about 500000 cluster centroids. At present the performance of code is okay when arr has tens of millions of samples and number of cluster centroids is small (10-20). But is very slow when the number of clusters is very large (hundreds of thousands).
How can I improve the runtime performance of this code for large number of cluster centers?.
First, some Python improvements. You do repeated lookups of the kind cluster_dict[cluster]. If you were to directly iterate over cluster_dict.values() these would all be not needed. Also, list comprehensions are usually faster than building a simple list with a for loop. And finally, docstrings go inside the function:
def update_center(self, cluster_dict):
    &quot;&quot;&quot;Gives new centroids with each new iteration.&quot;&quot;&quot;
    cluster_center = []

    for cluster in cluster_dict.values():
        cluster_list = [self.arr[sample] for sample in cluster]

        # Mean of samples in a given cluster.
        cluster_mean = np.mean(np.array(cluster_list), axis=0).tolist()

        # Find which sample is closest to mean of cluster.
        smallest_distance_from_mean = np.argmin([self.distance(cluster_mean, sample)
                                                 for sample in cluster_list])

        # Build an updated list of cluster centers.
        cluster_center.append(cluster[smallest_distance_from_mean])

    return cluster_center

Another small improvement is not to calculate the sqrt in your distance function, since you only ever need the minimum value, so the squared values work just fine.
def distance(self, arr1, arr2):
    &quot;&quot;&quot;Squared distance between two arrays&quot;&quot;&quot;
    return np.sum(np.subtract(arr1, arr2) ** 2)

However, the greatest improvements are probably to be gained by not using lists everywhere, but to use numpy wherever you can:
class KMeans:
    def __init__(self, arr, k):
        self.arr = np.array(arr)
        self.k = k

    def loop_over_centers(self, max_iter):

        # Choose random centroids.
        chosen_centers = np.random.choice(len(self.arr), self.k, replace=False)

        # Iterate till a certain max number of iterations.
        for _ in range(max_iter):

            cluster_dict = defaultdict(list)

            # Assign samples to cluster centers based on minimum distance
            centers = self.arr[chosen_centers]
            for i, sample in enumerate(self.arr):
                closest_center = np.argmin(np.sum((centers - sample)**2, axis=1))
                cluster_dict[chosen_centers[closest_center]].append(i)
            chosen_centers = self.update_center(cluster_dict)

        return cluster_dict

    def update_center(self, cluster_dict):
        &quot;&quot;&quot;Gives new centroids with each new iteration.&quot;&quot;&quot;
        cluster_center = []
    
        for cluster in cluster_dict.values():
            samples = self.arr[cluster]
            # Mean of samples in a given cluster.
            cluster_mean = np.mean(samples, axis=0)
            closest_sample = np.argmin(np.sum((samples - cluster_mean)**2, axis=1))
            cluster_center.append(cluster[closest_sample])

        return cluster_center

Your code takes 391 ms ± 21.7 ms on my machine, the first two changes drop this down to 340 ms ± 15.9 ms and the final change to 145 µs ± 588 ns, faster by a factor 2700!
You could go even further by doing this part without the for loop using only array operations as well:
        for i, sample in enumerate(self.arr):
            closest_center = np.argmin(np.sum((centers - sample)**2, axis=1))
            cluster_dict[chosen_centers[closest_center]].append(i)

However, this will then take k * len(arr) * dim(arr) memory, which may be a bit too much if both are very large.
Note: You can factor out the distance into a function again, I just don't want to rerun the timings...