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

原标题：缺乏 20 行 Python 代码，高效完成 k-means 均值聚类算法！

作者 | 许文武

责编 | 郭芮

出品 | CSDN 博客

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

关于 k-means 均值聚类算法的原理介绍、完成代码，网上有许多，但运转功率好像都有点问题。今日略微有点闲暇，写了一个缺乏20行的 k-means 均值聚类算法，1万个样本均匀耗时20毫秒（10次均值）。相同的数据样本，网上盛行的算法均匀耗时3000毫秒（10次均值）。间隔居然达百倍以上，令我深感意外，忍不住再次向 numpy 献上膝盖！

以下是我的代码，包括注释、空行一共26行，有用代码16行。

1importnumpy asnp

2

3defkmeans_xufive(ds, k):

4"""k-means聚类算法

5

6k - 指定分簇数量

7ds - ndarray(m, n)，m个样本的数据集，每个样本n个特点值

8"""

9

10m, n = ds.shape # m：样本数量，n：每个样本的特点值个数

11result = np.empty(m, dtype=np.int) # m个样本的聚类成果

12cores = np.empty((k, n)) # k个质心

13cores = ds[np.random.choice(np.arange(m), k, replace= False)] # 从m个数据样本中不重复地随机挑选k个样本作为质心

14

15whileTrue: # 迭代核算

16d = np.square(np.repeat(ds, k, axis= 0).reshape(m, k, n) - cores)

17distance = np.sqrt(np.sum(d, axis= 2)) # ndarray(m, k)，每个样本间隔k个质心的间隔，共有m行

18index_min = np.argmin(distance, axis= 1) # 每个样本间隔最近的质心索引序号

19

20if(index_min == result).all: # 假如样本聚类没有改动

21returnresult, cores # 则回来聚类成果和质心数据

22

23result[:] = index_min # 从头分类

24fori inrange(k): # 遍历质心集

25items = ds[result==i] # 找出对应当时质心的子样本集

26cores[i] = np.mean(items, axis= 0) # 以子样本集的均值作为当时质心的方位

这是网上比较盛行的 k-means 均值聚类算法代码，包括注释、空行一共57行，有用代码37行。

1importnumpy asnp

2

3# 加载数据

4defloadDataSet(fileName):

5data = np.loadtxt(fileName,delimiter= 't')

6returndata

7

8# 欧氏间隔核算

9defdistEclud(x,y):

10returnnp.sqrt(np.sum((x-y)** 2)) # 核算欧氏间隔

11

12# 为给定数据集构建一个包括K个随机质心的调集

13defrandCent(dataSet,k):

14m,n = dataSet.shape

15centroids = np.zeros((k,n))

16fori inrange(k):

17index = int(np.random.uniform( 0,m)) #

18centroids[i,:] = dataSet[index,:]

19returncentroids

20

21# k均值聚类

22defkmeans_open(dataSet,k):

23

24m = np.shape(dataSet)[ 0] #行的数目

25# 榜首列存样本属于哪一簇

26# 第二列存样本的到簇的中心点的差错

27clusterAssment = np.mat(np.zeros((m, 2)))

28clusterChange = True

29

30# 第1步 初始化centroids

31centroids = randCent(dataSet,k)

32whileclusterChange:

33clusterChange = False

34

35# 遍历一切的样本（行数）

36fori inrange(m):

37minDist = 100000.0

38minIndex = -1

39

40# 遍历一切的质心

41#第2步 找出最近的质心

42forj inrange(k):

43# 核算该样本到质心的欧式间隔

44distance = distEclud(centroids[j,:],dataSet[i,:])

45ifdistance < minDist:

46minDist = distance

47minIndex = j

48# 第 3 步：更新每一行样本所属的簇

49ifclusterAssment[i, 0] != minIndex:

50clusterChange = True

51clusterAssment[i,:] = minIndex,minDist** 2

52#第 4 步：更新质心

53forj inrange(k):

54pointsInCluster = dataSet[np.nonzero(clusterAssment[:, 0].A == j)[ 0]] # 获取簇类一切的点

55centroids[j,:] = np.mean(pointsInCluster,axis= 0) # 对矩阵的行求均值

56

57returnclusterAssment.A[:, 0], centroids

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

1defcreate_data_set(*cores):

2"""生成k-means聚类测验用数据集"""

3

4ds = list

5forx0, y0, z0 incores:

6x = np.random.normal(x0, 0.1+np.random.random/ 3, z0)

7y = np.random.normal(y0, 0.1+np.random.random/ 3, z0)

8ds.append(np.stack((x,y), axis= 1))

9

10returnnp.vstack(ds)

测验代码如下：

1importtime

2importmatplotlib.pyplot asplt

3

4k = 4

5ds = create_data_set(( 0, 0, 2500), ( 0, 2, 2500), ( 2, 0, 2500), ( 2, 2, 2500))

6

7t0 = time.time

8result, cores = kmeans_xufive(ds, k)

9t = time.time - t0

10

11plt.scatter(ds[:, 0], ds[:, 1], s= 1, c=result.astype(np.int))

12plt.scatter(cores[:, 0], cores[:, 1], marker= 'x', c=np.arange(k))

13plt.show

14

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

16

17t0 = time.time

18result, cores = kmeans_open(ds, k)

19t = time.time - t0

20

21plt.scatter(ds[:, 0], ds[:, 1], s= 1, c=result.astype(np.int))

22plt.scatter(cores[:, 0], cores[:, 1], marker= 'x', c=np.arange(k))

23plt.show

24

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

测验成果如下：

1PSD: XufiveGitCSDNcode> py-3. k-means.py

2运用 kmeans_xufive算法，1万个样本点，耗时0 .0156550.3秒

3运用 kmeans_open算法，1万个样本点，耗时3 .9990890.3秒

作用如下：

作者：许文武，博客昵称「天元浪子」，本文首发于作者CSDN博客https://blog.csdn.net/xufive/article/details/101448969。

责任编辑：