Python递归实现K均值算法

import random
import numpy as np

def kMeans(data, k, centroids=None, max_iter=100):
    """
    递归实现K均值算法
    :param data: 数据集，格式为数组或二维矩阵
    :param k: 类别数目
    :param centroids: 质心
    :param max_iter: 最大迭代次数
    :return: 分类结果、质心
    """
    # 随机初始化质心
    if centroids is None:
        centroids = random.sample(list(data), k)

    # 距离矩阵
    distance_matrix = np.zeros([len(data), k])

    # 分配簇
    clusters = [[] for _ in range(k)]
    for i, point in enumerate(data):
        for j, centroid in enumerate(centroids):
            distance_matrix[i][j] = np.linalg.norm(point - centroid)
        index = np.argmin(distance_matrix[i])
        clusters[index].append(i)

    # 计算新的质心
    new_centroids = np.zeros([k, len(data[0])])
    for i, cluster in enumerate(clusters):
        if len(cluster) != 0:
            new_centroids[i] = np.mean(data[cluster], axis=0)
        else:
            new_centroids[i] = centroids[i]

    # 递归停止条件
    if np.allclose(centroids, new_centroids, atol=1e-4) or max_iter == 0:
        labels = np.zeros(len(data))
        for i, cluster in enumerate(clusters):
            for j in cluster:
                labels[j] = i
        return labels, centroids
    else:
        return kMeans(data, k, new_centroids, max_iter-1)

# 测试代码
data = np.array([[1, 2], [2, 1], [3, 2], [6, 5], [7, 5], [8, 7]])
labels, centroids = kMeans(data, k=2)
print(labels)        # [0. 0. 0. 1. 1. 1.]
print(centroids)     # [[2. 1.] [7. 6.]]