使用Numpy实现K均值（K=2）聚类时的质心更新错误排查与代码优化求助

Let's break down the problems in your code step by step and fix them, plus add some optimizations to make it more efficient.

Key Issues in Your Current Code #

• Incorrect loop iteration: You used for i in iterations: where iterations is an integer (30). You can't iterate over an integer directly—replace this with for _ in range(iterations).
• Undefined variable reference: Inside the kmeans function, you referenced y in for c in range(len(y)), but y isn't defined in the function scope. Use len(centroids) instead, since that's the parameter holding your initial centroids.
• Non-numpy array inputs: Your input x and initial centroids are regular Python lists, not numpy arrays. Boolean indexing (like x[assignment == c]) won't work with regular lists, leading to the index error you saw.
• Unhandled empty clusters: If a centroid ends up with no assigned samples (unlikely in your small test case but possible in larger datasets), np.mean() will return nan—we should add a check to retain the original centroid in this scenario.

Fixed & Optimized Code #

First, let's optimize the Euclidean distance calculation—nested loops are inefficient for larger datasets, so we'll use numpy's broadcasting for vectorized computation:

import numpy as np

def euclidean_distance(x: np.ndarray, y: np.ndarray):
    # x shape: (N1, D)
    # y shape: (N2, D)
    # output shape: (N1, N2)
    # Vectorized calculation (far faster than nested loops)
    return np.sqrt(np.sum((x[:, np.newaxis] - y) ** 2, axis=2))

def kmeans(x, centroids, iterations=30):
    # Convert inputs to numpy arrays automatically if needed
    x = np.asarray(x)
    centroids = np.asarray(centroids)
    
    for _ in range(iterations):
        # Calculate distance matrix
        dist = euclidean_distance(x, centroids)
        # Assign each sample to the closest centroid
        assignment = np.argmin(dist, axis=1)
        
        # Update centroids (make a copy to avoid overwriting during iteration)
        new_centroids = centroids.copy()
        for c in range(len(new_centroids)):
            assigned_samples = x[assignment == c]
            # Only update centroid if there are assigned samples
            if len(assigned_samples) > 0:
                new_centroids[c] = np.mean(assigned_samples, axis=0)
        
        # Early exit if centroids stop changing (convergence)
        if np.allclose(centroids, new_centroids):
            break
        centroids = new_centroids
    
    return centroids, assignment

Testing the Fixed Code #

Run it with your test data:

x = [[1., 0.], [0., 1.], [0.5, 0.5]]
initial_centroids = [[1., 0.], [0., 1.]]

final_centroids, assignments = kmeans(x, initial_centroids)
print("Final Centroids:\n", final_centroids)
print("Sample Assignments:", assignments)

Sample Output #

Final Centroids:
 [[0.75 0.25]
 [0.   1.  ]]
Sample Assignments: [0 1 0]

(Note: The third sample is equidistant to initial centroids, so np.argmin assigns it to the first centroid index. The algorithm converges after one iteration here.)

Additional Improvements #

• Vectorization: The optimized distance function leverages numpy's broadcasting to avoid slow loops, making it scalable for larger datasets.
• Convergence Check: The np.allclose check lets the algorithm stop early once centroids stabilize, saving unnecessary iterations.
• Input Robustness: np.asarray ensures your code works whether you pass lists or numpy arrays as input.

内容的提问来源于stack exchange，提问作者Ha An