k-means算法代码实现

以下是Python中k-means算法的代码实现：

import numpy as np

class KMeans:
    def __init__(self, k=3, max_iter=100):
        self.k = k
        self.max_iter = max_iter
    
    def fit(self, X):
        # 随机初始化k个质心
        centroids = X[np.random.choice(X.shape[0], self.k, replace=False), :]
        
        for i in range(self.max_iter):
            # 分配每个样本点到最近的质心
            distances = np.sqrt(((X - centroids[:, np.newaxis])**2).sum(axis=2))
            labels = np.argmin(distances, axis=0)
            
            # 更新质心
            new_centroids = np.zeros((self.k, X.shape[1]))
            for j in range(self.k):
                new_centroids[j, :] = X[labels == j, :].mean(axis=0)
                
            # 如果质心不再发生变化则停止迭代
            if np.allclose(centroids, new_centroids):
                break
            else:
                centroids = new_centroids
                
        self.centroids = centroids
        self.labels = labels
        
    def predict(self, X):
        distances = np.sqrt(((X - self.centroids[:, np.newaxis])**2).sum(axis=2))
        labels = np.argmin(distances, axis=0)
        return labels

使用方法如下：

# 生成随机数据
X = np.random.randn(100, 2)

# 创建KMeans对象并进行训练
kmeans = KMeans(k=3)
kmeans.fit(X)

# 获取聚类结果
labels = kmeans.labels

# 获取质心
centroids = kmeans.centroids