kmeans源码python

以下是一个简单的 Python 实现 k-means 算法的代码。该代码使用 numpy 和 matplotlib 库进行数据处理和可视化。

import numpy as np
import matplotlib.pyplot as plt

def kmeans(X, k, max_iters=100):
    """
    K-means 聚类算法
    :param X: 数据集，ndarray 类型，每行表示一个样本
    :param k: 类别数
    :param max_iters: 最大迭代次数
    :return: centroids, clusters
    """
    m, n = X.shape  # 数据集大小
    centroids = X[np.random.choice(m, k, replace=False), :]  # 随机初始化质心
    clusters = np.zeros(m)  # 保存每个样本所属的类别
    
    for i in range(max_iters):
        # 计算每个样本到质心的距离
        distances = np.sqrt(np.sum((X - centroids[:, np.newaxis])**2, axis=2))
        
        # 分配类别
        clusters = np.argmin(distances, axis=0)
        
        # 更新质心
        for j in range(k):
            centroids[j] = np.mean(X[clusters == j], axis=0)
    
    return centroids, clusters

# 生成数据集并可视化
np.random.seed(42)
X = np.vstack([np.random.randn(30, 2) + [2, 2],
               np.random.randn(30, 2) + [0, -2],
               np.random.randn(30, 2) + [-2, 2]])
plt.scatter(X[:, 0], X[:, 1], s=50)
plt.show()

# 聚类并可视化
centroids, clusters = kmeans(X, k=3)
colors = ['r', 'g', 'b']
for i in range(3):
    plt.scatter(X[clusters == i, 0], X[clusters == i, 1], s=50, c=colors[i])
plt.scatter(centroids[:, 0], centroids[:, 1], marker='*', s=200, c='k')
plt.show()

在上面的代码中，我们首先生成了一个二维数据集，然后使用 kmeans 函数进行聚类，并将聚类结果可视化。你可以通过设置 k 值和 max_iters 参数来调整聚类的类别数和最大迭代次数。