Efficient Python Implementation of K-Means Clustering with Performance Comparison

The author explains that machine learning can be divided into four areas—classification, clustering, regression, and dimensionality reduction—and that k‑means, while simple, is an excellent introductory exercise.

A custom k‑means implementation written in fewer than 20 lines processes 10,000 samples in about 20 ms for ten iterations, which is over a hundred times faster than many popular versions.

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 = np.empty((k, n)) # k个质心 cores = ds[np.random.choice(np.arange(m), k, replace=False)] # 随机选择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个质心的距离 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) # 更新质心位置

The article also presents a more conventional k‑means version (kmeans_open) that includes helper functions for loading data, computing Euclidean distance, initializing random centroids, and the full clustering loop.

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)) 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 def kmeans_open(dataSet, k): m = np.shape(dataSet)[0] clusterAssment = np.mat(np.zeros((m, 2))) clusterChange = True centroids = randCent(dataSet, k) while clusterChange: clusterChange = False for i in range(m): minDist = 1e5 minIndex = -1 for j in range(k): distance = distEclud(centroids[j, :], dataSet[i, :]) if distance < minDist: minDist = distance minIndex = j if clusterAssment[i, 0] != minIndex: clusterChange = True clusterAssment[i, :] = minIndex, minDist ** 2 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

A utility function create_data_set generates synthetic data for testing, allowing the user to specify multiple centroids and the number of points around each.

def create_data_set(*cores): """生成k-means聚类测试用数据集""" ds = [] 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)

The testing script creates a dataset with four clusters (2,500 points each), runs both implementations, measures execution time, and visualizes the clustering results with Matplotlib.

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)) # Test custom implementation 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('custom kmeans_xufive: %.3f seconds' % t) # Test conventional implementation 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('conventional kmeans_open: %.3f seconds' % t)

Running the tests shows the custom version completes in roughly 0.016 seconds, while the conventional version takes about 4 seconds for the same 10,000‑point dataset, demonstrating a substantial speed advantage.