写一段k-means聚类的python代码

以下是一个简单的k-means聚类的Python代码：

import numpy as np

class KMeans:
    def __init__(self, k=2, max_iterations=100):
        self.k = k
        self.max_iterations = max_iterations

    def fit(self, X):
        self.centroids = X[np.random.choice(X.shape[0], self.k, replace=False)]
        for i in range(self.max_iterations):
            distances = np.sqrt(((X - self.centroids[:, np.newaxis])**2).sum(axis=2))
            cluster_index = np.argmin(distances, axis=0)
            for j in range(self.k):
                self.centroids[j] = np.mean(X[cluster_index == j], axis=0)

    def predict(self, X):
        distances = np.sqrt(((X - self.centroids[:, np.newaxis])**2).sum(axis=2))
        return np.argmin(distances, axis=0)

上述代码中，KMeans类的构造函数接受两个参数：k表示聚类的数量，max_iterations表示最大迭代次数。fit方法接受一个数据集X，初始化聚类中心，然后迭代更新聚类中心，直到达到最大迭代次数。predict方法接受一个数据集X，并返回每个样本所属的聚类编号。