KMeans聚类

原理and代码

"""
这里使用numpy来简单的手搓一下k-means的原理代码
"""
import numpy as np


def kmeans(X, K, max_iters=300):
    # 随机初始化 k 个聚类中心, 这里choice中的a只能是整数或者一维数组; 多维数组无法确定抽取什么
    # X: M * N， 其中N为特征数
    # centroids: K * N -> K * 1 * N
    # 想要得到一个: M * K * N 的矩阵, 需要为X增加一个维度, 变成: M * 1 * N
    centroids = X[np.random.choice(X.shape[0], K, replace=False)]
    
    for _ in range(max_iters):
        # distances = np.linalg.norm(X - centroids[:, np.newaxis], axis=2)
        distances = np.sqrt(((X[:, np.newaxis] - centroids) ** 2).sum(axis=2))
        # 算出每一个点属于哪一个聚类中心
        labels = np.argmin(distances, axis=1)
        # 计算新的聚类中心
        new_centroids = np.array([X[labels == k].mean(axis=0) for k in range(K)])
        # 如果聚类中心不再变化，停止迭代
        if np.all(new_centroids == centroids):
            break
        centroids = new_centroids

    return labels, centroids
    


if __name__ == "__main__":
    data = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
    kmeans(data, 2)