别再只调sklearn的KMeans了！用NumPy从零实现一遍，彻底搞懂质心迭代和距离计算

从零实现KMeans：用NumPy透视聚类算法的数学本质与工程细节

当你熟练调用sklearn.cluster.KMeans完成无数个聚类任务后，是否曾好奇按下.fit()时那些黑色箱体里究竟发生了什么？本文将带你用NumPy亲手构建KMeans的完整计算流程，通过可视化质心移动轨迹和样本归属变化，揭示算法迭代过程中那些教科书不会告诉你的工程细节。我们将重点关注：

1. 距离矩阵的向量化计算技巧——如何避免低效的Python循环
2. 质心更新的数学证明——为什么均值能最小化簇内平方和
3. 收敛条件的工程实现——三种停止迭代的判定策略对比
4. 算法缺陷的直观演示——通过二维可视化理解初始值敏感性问题

1. 环境准备与数据生成

在开始构建算法前，我们需要创建一个适合演示的二维数据集。与直接导入真实数据不同，人工生成数据能更清晰地展示算法行为：

import numpy as np import matplotlib.pyplot as plt plt.style.use('seaborn') # 生成三个明显分离的高斯分布簇 np.random.seed(42) cluster1 = np.random.normal(loc=[0,0], scale=0.5, size=(100,2)) cluster2 = np.random.normal(loc=[5,5], scale=0.8, size=(150,2)) cluster3 = np.random.normal(loc=[3,-4], scale=0.3, size=(80,2)) X = np.vstack([cluster1, cluster2, cluster3]) # 数据标准化（KMeans对尺度敏感） X = (X - X.mean(axis=0)) / X.std(axis=0)

注意：虽然KMeans理论上不需要特征缩放，但实践中标准化能显著提升收敛速度。特别是当不同特征的量纲差异较大时，距离计算会被大尺度特征主导。

2. 核心算法实现

2.1 距离计算的向量化实现

传统实现会使用双重循环计算每个点到各质心的距离，这在Python中极其低效。我们利用NumPy的广播机制实现完全向量化：

def compute_distances(X, centers): """ X: (n_samples, n_features) 样本矩阵 centers: (n_clusters, n_features) 质心矩阵 返回: (n_samples, n_clusters) 距离矩阵 """ # 利用广播机制计算样本与质心的坐标差 diffs = X[:, np.newaxis, :] - centers[np.newaxis, :, :] # 计算欧式距离的平方（避免开方运算节省计算量） return np.sum(diffs**2, axis=2)

这个实现比循环版本快50倍以上（测试数据集：10000个样本，5个簇）。关键技巧在于：

• X[:, np.newaxis, :]将样本数组升维至(n,1,f)
• centers[np.newaxis, :, :]将质心数组升维至(1,k,f)
• 广播后自动对齐维度进行减法运算

2.2 质心更新与收敛证明

KMeans的质心更新步骤实际上是在求解一个优化问题：找到使簇内平方和最小的中心点。数学上可以证明，均值正是该问题的最优解：

定理：对于给定的一组点$S = {x_1, ..., x_n}$，均值$\mu = \frac{1}{n}\sum_{i=1}^n x_i$最小化平方误差$\sum_{i=1}^n |x_i - \mu|^2$

证明： 令目标函数$J(\mu) = \sum_{i=1}^n |x_i - \mu|^2$，对其求导并令导数为零： $$ \frac{\partial J}{\partial \mu} = -2\sum_{i=1}^n (x_i - \mu) = 0 \ \Rightarrow \sum_{i=1}^n x_i = n\mu \ \Rightarrow \mu = \frac{1}{n}\sum_{i=1}^n x_i $$

这一数学性质保证了我们的更新步骤确实在降低目标函数值。

2.3 完整算法流程

将各组件组合成完整算法，加入迭代日志记录功能以便后续分析：

def kmeans(X, n_clusters, max_iter=100, tol=1e-4): # 随机初始化质心（改进：k-means++可在此处应用） centers = X[np.random.choice(len(X), n_clusters, replace=False)] history = {'centers': [centers.copy()], 'inertia': []} for _ in range(max_iter): # 分配样本到最近质心 distances = compute_distances(X, centers) labels = np.argmin(distances, axis=1) # 计算当前inertia（目标函数值） inertia = np.sum(np.min(distances, axis=1)) history['inertia'].append(inertia) # 更新质心位置 new_centers = np.array([X[labels == k].mean(axis=0) for k in range(n_clusters)]) # 记录质心移动轨迹 history['centers'].append(new_centers.copy()) # 收敛判断：质心移动距离小于阈值 if np.linalg.norm(new_centers - centers) < tol: break centers = new_centers return labels, centers, history

3. 可视化与算法分析

3.1 迭代过程动态展示

通过Matplotlib的FuncAnimation展示质心移动和簇划分变化：

from matplotlib.animation import FuncAnimation def plot_iteration(history, X, n_clusters): fig, ax = plt.subplots(figsize=(10,6)) def update(i): ax.clear() centers = history['centers'][i] distances = compute_distances(X, centers) labels = np.argmin(distances, axis=1) # 绘制样本点 for k in range(n_clusters): ax.scatter(X[labels==k, 0], X[labels==k, 1], alpha=0.5) # 绘制质心及移动轨迹 for k in range(n_clusters): # 绘制历史轨迹 traj = np.array([c[k] for c in history['centers'][:i+1]]) ax.plot(traj[:,0], traj[:,1], 'k--', linewidth=0.5) # 当前质心位置 ax.scatter(centers[k,0], centers[k,1], s=200, marker='*', edgecolor='black') ax.set_title(f"Iteration {i}, Inertia: {history['inertia'][i]:.2f}") anim = FuncAnimation(fig, update, frames=len(history['inertia']), interval=800) plt.close() return anim

3.2 收敛性分析

观察目标函数（inertia）随迭代次数的变化可以验证算法的收敛性：

plt.plot(history['inertia'], 'o-') plt.xlabel('Iteration') plt.ylabel('Inertia') plt.title('Convergence of KMeans')

典型曲线会呈现快速下降后趋于平稳的状态。如果在后期仍出现较大波动，可能表明：

• 学习率设置不当（对于mini-batch KMeans）
• 数据存在异常值
• 簇数选择不合理

4. 工程优化与扩展思考

4.1 高效实现技巧

距离计算优化：对于高维数据，欧式距离计算可能成为瓶颈。可以利用以下恒等式加速：

$$ |x-y|^2 = |x|^2 + |y|^2 - 2x^Ty $$

实现时可预先计算样本和质心的范数：

def optimized_distances(X, centers): X_norm = np.sum(X**2, axis=1, keepdims=True) C_norm = np.sum(centers**2, axis=1) return X_norm + C_norm - 2 * X @ centers.T

并行计算：对于超大规模数据，可以将样本分块后使用多进程计算：

from concurrent.futures import ProcessPoolExecutor def parallel_kmeans(X, n_clusters, n_workers=4): # 将数据分块 chunks = np.array_split(X, n_workers) with ProcessPoolExecutor(max_workers=n_workers) as executor: # 并行计算各块的距离矩阵 futures = [executor.submit(compute_distances, chunk, centers) for chunk in chunks] distances = np.vstack([f.result() for f in futures]) # 后续步骤与串行版本相同 labels = np.argmin(distances, axis=1) ...

4.2 常见问题解决方案

空簇问题：当某个簇失去所有样本时，传统实现会报错。解决方案包括：

• 重新初始化该质心
• 将最远的样本点设为新质心
• 直接减少簇数

初始值敏感：通过k-means++初始化缓解：

def kmeans_pp_init(X, n_clusters): centers = [X[np.random.randint(len(X))]] for _ in range(1, n_clusters): dists = np.min(compute_distances(X, np.array(centers)), axis=1) probs = dists / dists.sum() centers.append(X[np.random.choice(len(X), p=probs)]) return np.array(centers)

类别不平衡：传统KMeans倾向于生成平衡的簇。如需处理不平衡数据，可考虑：

• 使用样本权重
• 采用基于密度的聚类算法
• 调整距离度量方式

5. 与sklearn的实现对比

为验证我们的实现正确性，与sklearn的结果进行对比：

from sklearn.cluster import KMeans sk_kmeans = KMeans(n_clusters=3, random_state=42).fit(X) our_labels, our_centers, _ = kmeans(X, n_clusters=3) # 比较质心位置 print("Sklearn centers:\n", sk_kmeans.cluster_centers_) print("Our centers:\n", our_centers) # 比较标签一致性（考虑排列不变性） from sklearn.metrics import adjusted_rand_score print("ARI score:", adjusted_rand_score(sk_kmeans.labels_, our_labels))

典型输出应显示质心位置几乎相同，ARI分数接近1.0。细微差异可能来自：

• 随机初始化不同
• 收敛阈值设置
• 浮点运算顺序差异