不足20行 python 代码，高效实现 k-means 均值聚类算法

scikti-learn 将机器学习分为4个领域，分别是分类(classification)、聚类(clustering)、回归(regression)和降维(dimensionality reduction)。k-means均值算法虽然是聚类算法中比较简单的一种，却包含了丰富的思想内容，非常适合作为初学者的入门习题。

关于 k-means 均值聚类算法的原理介绍、实现代码，网上有很多，但运行效率似乎都有点问题。今天稍微有点空闲，写了一个不足20行的 k-means 均值聚类算法，1万个样本平均耗时20毫秒（10次均值）。同样的数据样本，网上流行的算法平均耗时3000毫秒（10次均值）。差距竟然达百倍以上，令我深感意外，不由得再次向 numpy 献上膝盖！

以下是我的代码，包含注释、空行总共25行，有效代码15行。

import numpy as np

def kmeans_xufive(ds, k):
    """k-means聚类算法
    
    k       - 指定分簇数量
    ds      - ndarray(m, n)，m个样本的数据集，每个样本n个属性值
    """
    
    m, n = ds.shape # m：样本数量，n：每个样本的属性值个数
    result = np.empty(m, dtype=np.int) # m个样本的聚类结果
    cores = ds[np.random.choice(np.arange(m), k, replace=False)] # 从m个数据样本中不重复地随机选择k个样本作为质心
    
    while True: # 迭代计算
        d = np.square(np.repeat(ds, k, axis=0).reshape(m, k, n) - cores)
        distance = np.sqrt(np.sum(d, axis=2)) # ndarray(m, k)，每个样本距离k个质心的距离，共有m行
        index_min = np.argmin(distance, axis=1) # 每个样本距离最近的质心索引序号
        
        if (index_min == result).all(): # 如果样本聚类没有改变
            return result, cores # 则返回聚类结果和质心数据
        
        result[:] = index_min # 重新分类
        for i in range(k): # 遍历质心集
            items = ds[result==i] # 找出对应当前质心的子样本集
            cores[i] = np.mean(items, axis=0) # 以子样本集的均值作为当前质心的位置

这是网上比较流行的 k-means 均值聚类算法代码，包含注释、空行总共57行，有效代码37行。

import numpy as np

# 加载数据
def loadDataSet(fileName):
    data = np.loadtxt(fileName,delimiter='\t')
    return data
 
# 欧氏距离计算
def distEclud(x,y):
    return np.sqrt(np.sum((x-y)**2))  # 计算欧氏距离
 
# 为给定数据集构建一个包含K个随机质心的集合
def randCent(dataSet,k):
    m,n = dataSet.shape
    centroids = np.zeros((k,n))
    for i in range(k):
        index = int(np.random.uniform(0,m)) #
        centroids[i,:] = dataSet[index,:]
    return centroids
 
# k均值聚类
def kmeans_open(dataSet,k):
 
    m = np.shape(dataSet)[0]  #行的数目
    # 第一列存样本属于哪一簇
    # 第二列存样本的到簇的中心点的误差
    clusterAssment = np.mat(np.zeros((m,2)))
    clusterChange = True
 
    # 第1步 初始化centroids
    centroids = randCent(dataSet,k)
    while clusterChange:
        clusterChange = False
 
        # 遍历所有的样本（行数）
        for i in range(m):
            minDist = 100000.0
            minIndex = -1
 
            # 遍历所有的质心
            #第2步 找出最近的质心
            for j in range(k):
                # 计算该样本到质心的欧式距离
                distance = distEclud(centroids[j,:],dataSet[i,:])
                if distance < minDist:
                    minDist = distance
                    minIndex = j
            # 第 3 步：更新每一行样本所属的簇
            if clusterAssment[i,0] != minIndex:
                clusterChange = True
                clusterAssment[i,:] = minIndex,minDist**2
        #第 4 步：更新质心
        for j in range(k):
            pointsInCluster = dataSet[np.nonzero(clusterAssment[:,0].A == j)[0]]  # 获取簇类所有的点
            centroids[j,:] = np.mean(pointsInCluster,axis=0)   # 对矩阵的行求均值
 
    return clusterAssment.A[:,0], centroids

函数create_data_set()，用于生成测试数据。可变参数 cores 是多个三元组，每一个三元组分别是质心的x坐标、y坐标和对应该质心的数据点的数量。

def create_data_set(*cores):
    """生成k-means聚类测试用数据集"""
    
    ds = list()
    for x0, y0, z0 in cores:
        x = np.random.normal(x0, 0.1+np.random.random()/3, z0)
        y = np.random.normal(y0, 0.1+np.random.random()/3, z0)
        ds.append(np.stack((x,y), axis=1))
    
    return np.vstack(ds)

测试代码如下:

import time
import matplotlib.pyplot as plt

k = 4
ds = create_data_set((0,0,2500), (0,2,2500), (2,0,2500), (2,2,2500))

t0 = time.time()
result, cores = kmeans_xufive(ds, k)
t = time.time() - t0

plt.scatter(ds[:,0], ds[:,1], s=1, c=result.astype(np.int))
plt.scatter(cores[:,0], cores[:,1], marker='x', c=np.arange(k))
plt.show()

print(u'使用kmeans_xufive算法，1万个样本点，耗时%f0.3秒'%t)

t0 = time.time()
result, cores = kmeans_open(ds, k)
t = time.time() - t0

plt.scatter(ds[:,0], ds[:,1], s=1, c=result.astype(np.int))
plt.scatter(cores[:,0], cores[:,1], marker='x', c=np.arange(k))
plt.show()

print(u'使用kmeans_open算法，1万个样本点，耗时%f0.3秒'%t)

测试结果如下：

PS D:\XufiveGit\CSDN\code> py -3 .\k-means.py
使用kmeans_xufive算法，1万个样本点，耗时0.0156550.3秒
使用kmeans_open算法，1万个样本点，耗时3.9990890.3秒

效果如下：

后记

近期有很多朋友通过私信咨询有关python学习问题。为便于交流，我在CSDN的app上创建了一个小组，名为“python作业辅导小组”，面向python初学者，为大家提供咨询服务、辅导python作业。欢迎有兴趣的同学扫码加入。

CSDN 不止为我们提供了这样一个交流平台，还经常推出各类技术交流活动。近期我将在 GeekTalk 栏目，和 Python 新手共同探讨如何快速成长为基础扎实、功力强大的程序员。CSDN 还为这个活动提供了一些纪念品。如果有兴趣，请扫码加入，或者直接点此进入。