如何并行化KMeans？

import numpy as np
from sklearn.cluster import KMeans
from multiprocessing import Pool

def kmeans_parallel(data, k, n_jobs):
    def kmeans_worker(data_chunk):
        kmeans = KMeans(n_clusters=k)
        kmeans.fit(data_chunk)
        return kmeans.cluster_centers_

    # Split data into chunks
    chunk_size = len(data) // n_jobs
    data_chunks = [data[i:i + chunk_size] for i in range(0, len(data), chunk_size)]

    # Run KMeans in parallel
    with Pool(n_jobs) as pool:
        results = pool.map(kmeans_worker, data_chunks)

    # Combine results
    new_centers = np.vstack(results)
    return new_centers

# Example usage
data = np.random.rand(1000, 10)
k = 3
n_jobs = 4
new_centers = kmeans_parallel(data, k, n_jobs)
print(new_centers)

import numpy as np
from sklearn.metrics import pairwise_distances_argmin_min

def update_centers_parallel(data, labels, k):
    centers = np.zeros((k, data.shape[1]))
    for i in range(k):
        points = data[labels == i]
        if len(points) > 0:
            centers[i] = np.mean(points, axis=0)
    return centers

def kmeans_parallel(data, k, max_iters=100, n_jobs=4):
    # Initialize centers
    centers = data[np.random.choice(data.shape[0], k, replace=False)]

    for _ in range(max_iters):
        # Assign points to nearest center
        labels = pairwise_distances_argmin_min(data, centers)[1]

        # Update centers in parallel
        new_centers = update_centers_parallel(data, labels, k)

        # Check for convergence
        if np.all(centers == new_centers):
            break
        centers = new_centers

    return centers, labels

# Example usage
data = np.random.rand(1000, 10)
k = 3
centers, labels = kmeans_parallel(data, k)
print(centers)