K-Means聚类算法python代码

以下是K-Means聚类算法的Python代码：

import numpy as np
import matplotlib.pyplot as plt

class KMeans:
    def __init__(self, k=3, max_iter=100):
        self.k = k
        self.max_iter = max_iter
        
    def fit(self, X):
        # 随机初始化质心
        self.centers = X[np.random.choice(X.shape[0], self.k, replace=False)]
        for i in range(self.max_iter):
            # 分配数据点到最近的质心
            labels = np.argmin(np.sqrt(((X - self.centers[:, np.newaxis])**2).sum(axis=2)), axis=0)
            # 更新质心
            new_centers = np.array([X[labels == j].mean(axis=0) for j in range(self.k)])
            # 如果质心不再移动，退出循环
            if np.allclose(self.centers, new_centers):
                break
            self.centers = new_centers
            
    def predict(self, X):
        return np.argmin(np.sqrt(((X - self.centers[:, np.newaxis])**2).sum(axis=2)), axis=0)

使用示例：

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

# 初始化KMeans算法
kmeans = KMeans(k=3)

# 训练模型
kmeans.fit(X)

# 预测数据
labels = kmeans.predict(X)

# 可视化结果
plt.scatter(X[:, 0], X[:, 1], c=labels)
plt.scatter(kmeans.centers[:, 0], kmeans.centers[:, 1], marker='x', s=200, linewidths=3, color='r')
plt.show()

输出结果如下所示：