纯Numpy手写K-Means：从原理到工程实践

1. 为什么还要手写 K-Means

很多人第一次接触聚类，都是从sklearn.cluster.KMeans开始的。两行代码，fit一下，labels_一取，任务就算完成了。但如果你真的在生产环境里调过参、排查过收敛问题，就会发现黑盒用起来爽，出问题的时候却无从下手。距离度量到底怎么算的、质心怎么初始化的、空簇怎么处理、迭代什么时候停——这些细节在封装好的接口里全被藏起来了。

用纯 Numpy 手写一遍 K-Means，最大的价值不是“造轮子”，而是把上面这些环节全部摊开在你自己面前。你会被迫去思考：为什么初始化方式对结果影响这么大？为什么同样的数据跑两次结果不一样？为什么有的簇会莫名其妙消失？这些问题在调包时你可能一辈子都遇不到，但一旦遇到，手写一遍的经历就是你的底气。

这篇文章适合三类人：一是刚学完 Numpy 想找个练手项目的初学者；二是用惯了 sklearn 但想搞清楚底层逻辑的工程师；三是需要在特殊场景（比如自定义距离、约束聚类）下改造算法的人。我会从设计思路讲到逐行实现，再把我自己踩过的坑和调优经验一并倒出来。代码全部基于 Numpy，不依赖任何机器学习框架，复制就能跑。

2. 整体设计与核心思路拆解

2.1 K-Means 到底在优化什么

K-Means 的目标非常朴素：把 N 个样本分成 K 个簇，使得每个样本到它所属簇质心的距离平方和最小。用公式写出来就是：

$$J = \sum_{k=1}^{K} \sum_{x_i \in C_k} |x_i - \mu_k|^2$$

其中 $\mu_k$ 是第 k 个簇的质心，$C_k$ 是第 k 个簇的样本集合。这个目标函数也叫 SSE（Sum of Squared Errors）或者惯性（inertia）。

理解这个公式是理解整个算法的钥匙。它告诉我们两件事：第一，距离用的是欧氏距离的平方，这决定了算法对量纲敏感，做之前必须标准化；第二，优化的是一个非凸目标，存在大量局部最优解，所以初始化方式直接决定了你最终落在哪个坑里。

算法本身是一个迭代过程，交替执行两个步骤：

• 分配步骤：固定质心，把每个样本分配给最近的质心
• 更新步骤：固定分配关系，把每个质心移动到所属样本的均值位置

这两步交替执行，SSE 单调不增，所以算法一定收敛。但收敛到的是局部最优，不是全局最优。这就是为什么实践中要跑多次取最优。

2.2 为什么选择纯 Numpy 实现

有人会问，既然 sklearn 已经写得那么好了，为什么还要自己写？我的理由有这么几条。

第一是可控性。sklearn 的 KMeans 默认用 k-means++ 初始化，跑 10 次取最优，这些默认值在大多数场景下没问题，但当你需要固定随机种子做可复现实验、或者需要自定义初始化策略时，就得改源码或者绕弯子。自己写的话，初始化策略就是一个函数，想换就换。

第二是可扩展性。标准 K-Means 用的是欧氏距离，但实际业务里你可能需要余弦距离、曼哈顿距离，甚至带权重的距离。sklearn 虽然支持 precomputed 距离矩阵，但内存开销大。自己写的话，距离计算就是一个函数调用，换成什么距离都行。

第三是教学价值。这一点对初学者尤其重要。手写一遍之后，你对“向量化”“广播机制”“收敛判据”这些概念的理解会上一个台阶。Numpy 的广播机制在 K-Means 里体现得淋漓尽致，写完这个项目，你对[:, None, :]这种写法的理解会深刻很多。

第四是性能可控。纯 Numpy 的向量化实现，在中等规模数据上（几万样本、几十维特征）性能其实不差，而且没有框架启动开销。如果你的场景是嵌入到一个轻量级服务里，不想引入 sklearn 这个依赖，手写版本反而更合适。

2.3 整体架构设计

我把整个实现拆成四个模块，每个模块职责单一，方便单独测试和替换：

这种拆分的好处是，每个模块都可以独立单元测试。比如初始化模块，我可以单独验证 k-means++ 选出来的点是不是真的分散；距离模块，我可以用手算的小矩阵验证广播写法是否正确。

3. 核心细节解析与实操要点

3.1 距离矩阵的向量化计算

这是整个实现里最容易写错、也最能体现 Numpy 功底的地方。假设样本矩阵X形状是(N, D)，质心矩阵C形状是(K, D)，我们要算一个(N, K)的距离矩阵。

最直观的写法是双重循环，但 N 一大就慢得没法看。向量化的写法利用了一个恒等式：

$$|x - c|^2 = |x|^2 - 2x \cdot c + |c|^2$$

对应到 Numpy：

def euclidean_distances(X, C): # X: (N, D), C: (K, D) X_sq = np.sum(X ** 2, axis=1, keepdims=True) # (N, 1) C_sq = np.sum(C ** 2, axis=1, keepdims=True).T # (1, K) cross = X @ C.T # (N, K) dist_sq = X_sq - 2 * cross + C_sq # (N, K) return np.maximum(dist_sq, 0) # 防止浮点误差导致负数

这里有几个细节值得说。keepdims=True是为了保持二维形状，方便后面的广播。np.maximum(dist_sq, 0)是必须的，因为浮点运算可能让理论上为 0 的值变成-1e-12，开根号就出 NaN 了。这个坑我在实际项目里踩过，当时数据里有一批重复点，距离算出来是负的极小值，sqrt之后全是 NaN，排查了半天才定位到。

如果你只需要分配结果，不需要真实距离值，那连开根号都可以省掉，直接比较平方距离就行。因为平方是单调函数，最近邻关系不变。这个优化在 K 很大或者 N 很大时能省不少时间。

3.2 初始化策略：随机 vs k-means++

初始化是 K-Means 里最关键的环节，没有之一。随机初始化的问题在于，如果两个初始质心恰好落在同一个密集簇里，那这个簇就会被劈成两半，而另一个真实存在的簇可能完全没有质心覆盖，最终结果惨不忍睹。

k-means++ 的思路很聪明：第一个质心随机选，之后的每个质心，按照“到已有质心距离平方”的概率分布来选。距离越远的点，被选中的概率越大。这样选出来的初始质心天然分散。

def kmeans_plus_plus_init(X, K, rng): N, D = X.shape centers = np.empty((K, D), dtype=X.dtype) # 第一个质心随机选 idx = rng.integers(N) centers[0] = X[idx] # 记录每个点到最近质心的距离平方 closest_dist_sq = euclidean_distances(X, centers[:1]).ravel() for k in range(1, K): # 按距离平方做加权采样 probs = closest_dist_sq / closest_dist_sq.sum() idx = rng.choice(N, p=probs) centers[k] = X[idx] # 更新最近距离 new_dist_sq = euclidean_distances(X, centers[k:k+1]).ravel() closest_dist_sq = np.minimum(closest_dist_sq, new_dist_sq) return centers

这里rng.choice(N, p=probs)是核心，它实现了按概率采样。注意probs必须归一化，而且不能有负数。如果数据里有重复点导致closest_dist_sq全为 0，那probs就是 NaN，需要加一个极小值兜底。

实测下来，k-means++ 在大多数数据集上能把最终 SSE 降低 20% 到 50%，代价是初始化阶段多花一点时间。对于 K 不大的场景，这点开销完全值得。

3.3 空簇处理：一个容易被忽略的边界情况

迭代过程中，某个簇可能一个样本都没分配到，这时候它的质心就是空的，均值计算会出问题。sklearn 的处理方式是把最远的点重新分配给空簇，但具体实现细节没有公开。

我的处理策略是：如果发现空簇，就把距离它当前质心最远的那个样本点“抢”过来，作为这个空簇的新质心。这样既保证了 K 个簇都存在，又不会让空簇的质心乱飘。

def handle_empty_clusters(X, labels, centers, dist_matrix): K = centers.shape[0] for k in range(K): if np.sum(labels == k) == 0: # 找到当前分配中距离自己质心最远的样本 assigned_dist = dist_matrix[np.arange(len(X)), labels] far_idx = np.argmax(assigned_dist) labels[far_idx] = k centers[k] = X[far_idx] return labels, centers

这个逻辑要放在分配步骤之后、更新步骤之前。注意dist_matrix是分配时算好的距离矩阵，直接复用可以省一次计算。

提示：空簇处理不是必须的，如果你的数据分布比较均匀、K 选得合理，空簇很少出现。但一旦出现而不处理，程序就会报错或者结果错乱。生产环境里建议加上这个保护。

3.4 收敛判据的选择

收敛判据有两种常见写法：一种是看质心移动距离，一种是看标签变化比例。

质心移动距离的写法是：如果所有质心移动的最大距离小于tol，就认为收敛。这个判据比较严格，因为质心可能在小范围内震荡。

标签变化的写法是：如果这一轮和上一轮相比，标签变化的样本比例小于某个阈值，就认为收敛。这个判据更贴近业务含义，因为最终我们要的是标签。

我一般两个都用，取“或”的关系：质心移动足够小，或者标签基本不变，都算收敛。这样能避免不必要的迭代。

def has_converged(old_centers, new_centers, old_labels, new_labels, tol=1e-4): center_shift = np.max(np.linalg.norm(new_centers - old_centers, axis=1)) label_change = np.mean(old_labels != new_labels) return center_shift < tol or label_change < tol

tol的取值需要根据数据量纲来定。如果数据做过标准化，1e-4是个合理的默认值。如果没标准化，质心移动距离可能很大，这个阈值就不适用了。

4. 完整实操过程与核心环节实现

4.1 数据准备与标准化

K-Means 对量纲极度敏感。假设你有两个特征，一个是年龄（0-100），一个是年收入（0-1000000），那距离计算完全被收入主导，年龄特征形同虚设。所以标准化是必须的前置步骤。

def standardize(X): mean = X.mean(axis=0) std = X.std(axis=0) std[std == 0] = 1.0 # 防止除零 return (X - mean) / std, mean, std

注意std == 0的处理。如果某个特征所有值都一样，标准差是 0，除零会出 NaN。这种情况在真实数据里不少见，比如某个字段全是默认值。

标准化之后，所有特征均值为 0、方差为 1，距离计算才公平。如果你需要把结果还原到原始尺度，保存mean和std就行。

4.2 主循环的完整实现

把前面的模块串起来，主循环长这样：

class NumpyKMeans: def __init__(self, n_clusters=8, max_iter=300, tol=1e-4, init='k-means++', n_init=10, random_state=None): self.n_clusters = n_clusters self.max_iter = max_iter self.tol = tol self.init = init self.n_init = n_init self.random_state = random_state def fit(self, X): rng = np.random.default_rng(self.random_state) best_inertia = np.inf best_labels = None best_centers = None for _ in range(self.n_init): centers = self._init_centers(X, rng) labels = np.zeros(len(X), dtype=int) for i in range(self.max_iter): dist = euclidean_distances(X, centers) new_labels = np.argmin(dist, axis=1) # 空簇处理 new_labels, centers = handle_empty_clusters( X, new_labels, centers, dist) # 更新质心 new_centers = np.array([ X[new_labels == k].mean(axis=0) if np.sum(new_labels == k) > 0 else centers[k] for k in range(self.n_clusters) ]) if has_converged(centers, new_centers, labels, new_labels, self.tol): labels = new_labels centers = new_centers break labels = new_labels centers = new_centers inertia = self._compute_inertia(X, labels, centers) if inertia < best_inertia: best_inertia = inertia best_labels = labels.copy() best_centers = centers.copy() self.labels_ = best_labels self.cluster_centers_ = best_centers self.inertia_ = best_inertia return self def _init_centers(self, X, rng): if self.init == 'k-means++': return kmeans_plus_plus_init(X, self.n_clusters, rng) else: idx = rng.choice(len(X), self.n_clusters, replace=False) return X[idx].copy() def _compute_inertia(self, X, labels, centers): dist = euclidean_distances(X, centers) return np.sum(dist[np.arange(len(X)), labels])

n_init=10是跑 10 次不同初始化，取 SSE 最小的那次。这个策略在 sklearn 里是默认的，实测下来能显著提升结果稳定性。代价是时间乘以 10，如果数据量很大，可以适当降低。

4.3 用模拟数据验证

我构造了三簇二维高斯数据来验证实现：

rng = np.random.default_rng(42) X1 = rng.normal(loc=[0, 0], scale=0.5, size=(200, 2)) X2 = rng.normal(loc=[5, 5], scale=0.5, size=(200, 2)) X3 = rng.normal(loc=[0, 5], scale=0.5, size=(200, 2)) X = np.vstack([X1, X2, X3]) model = NumpyKMeans(n_clusters=3, random_state=42) model.fit(X) print(f"SSE: {model.inertia_:.2f}") print(f"质心:\n{model.cluster_centers_}")

跑出来的质心应该接近[0,0]、[5,5]、[0,5]这三个真实中心。SSE 大概在 300 左右（因为每簇方差是 0.5，200 个点，理论 SSE 约 3 * 200 * 2 * 0.25 = 300）。

如果结果偏离很大，先检查标准化有没有做，再检查 k-means++ 的采样概率有没有归一化。

4.4 性能优化：从 O(NKD) 到更快

标准实现每轮迭代要算一次完整的距离矩阵，复杂度是 O(NKD)。当 N=100万、K=100、D=50 时，一次距离计算就是 50 亿次浮点运算，很吃力。

几个优化方向：

第一，用平方距离代替真实距离。前面说过，比较最近邻不需要开根号，省掉sqrt能快 20% 左右。

第二，利用三角不等式剪枝。如果某个样本到当前质心的距离已经小于到其他质心距离的一半，那它不可能换簇，可以跳过计算。Elkan 算法就是基于这个思路，能快好几倍，但实现复杂。

第三，Mini-Batch K-Means。每次只用一小批样本更新质心，大幅降低单轮计算量。代价是结果略差，但速度快一个数量级。sklearn 的MiniBatchKMeans就是这个思路。

第四，用 float32 代替 float64。内存减半，速度提升 30% 左右，精度损失在聚类场景下可以忽略。

我一般先用 float32 + 平方距离，如果还不够快再考虑 Mini-Batch。Elkan 算法除非数据量特别大，否则不值得那个实现复杂度。

5. 常见问题与排查技巧实录

5.1 结果每次都不一样怎么办

这是新手最常问的问题。原因有两个：一是随机种子没固定，二是n_init太小。

固定种子很简单，random_state=42传进去就行。但要注意，Numpy 的新版随机数生成器default_rng和旧版np.random.seed是两套体系，混用会导致种子失效。统一用default_rng就没问题。

n_init建议至少 10。如果数据簇结构不明显，可以加到 20 甚至 50。代价是时间线性增长，但换来的是结果稳定。

还有一个隐藏原因：如果数据里有重复点，k-means++ 采样时概率可能全为 0，导致选出来的质心都一样。这时候需要给概率加一个极小值兜底。

5.2 簇的大小严重不均

有时候跑出来一个簇有 9000 个点，另外两个簇各 500 个。这通常不是 bug，而是数据本身就不均衡，或者 K 选大了。

K-Means 有个固有偏好：它倾向于产生大小相近的球形簇。如果真实簇大小差异很大，K-Means 会把大簇劈开，或者把小簇合并。这不是算法错了，是模型假设和数据不匹配。

解决办法有几个：一是换算法，比如 DBSCAN 或者高斯混合模型；二是对样本加权，让大簇的样本权重降低；三是用 K-Medoids，它对异常值更鲁棒。

5.3 SSE 下降很慢或者震荡

正常情况下 SSE 应该单调下降。如果出现震荡，说明代码有 bug。最常见的原因是空簇处理逻辑写错了，导致质心在两步之间来回跳。

排查方法：打印每轮的 SSE 和质心移动距离。如果 SSE 上升，检查更新步骤是不是用了旧的标签。如果质心移动距离一直很大但不收敛，检查tol是不是设得太小。

还有一种可能是数据没标准化，导致距离计算数值不稳定。标准化之后一般就好了。

5.4 如何选择 K 值

这是聚类里最玄学的问题。常用的方法有肘部法、轮廓系数、Gap Statistic。

肘部法是画 SSE 随 K 变化的曲线，找拐点。但拐点往往不明显，主观性很强。

轮廓系数更客观，计算每个点的轮廓系数（它到同簇其他点的平均距离 vs 到最近其他簇的平均距离），取平均。值域 [-1, 1]，越大越好。缺点是计算复杂度 O(N^2)，大数据集跑不动。

我的经验是：先用肘部法粗筛，再用轮廓系数在小范围内精选。如果业务上有明确的簇数预期，直接按业务来，别迷信指标。

5.5 几个我踩过的坑

第一个坑是整数除法。早期我用 Python 2 写的时候，np.sum(labels == k) / len(X)结果是 0，因为整数除法截断了。现在 Python 3 没这个问题，但如果你从别处拷贝代码，要注意。

第二个坑是广播维度对不上。X - centers这种写法，如果X是(N, D)、centers是(K, D)，直接减会报错。必须写成X[:, None, :] - centers[None, :, :]，得到(N, K, D)的中间结果。这个中间结果内存开销是 NKD，大数据集下会爆内存。所以我才用恒等式展开的写法，避免生成三维数组。

第三个坑是空簇导致的 NaN。如果某个簇没有样本，X[labels == k].mean(axis=0)会返回 NaN 并报警告。必须加判断。

第四个坑是收敛判据用错。我一开始只看质心移动距离，结果在某些数据上质心几乎不动但标签还在变，导致提前停止。后来改成两个判据取或，就稳定了。

6. 扩展方向与个人体会

手写 K-Means 之后，往几个方向扩展都很自然。

自定义距离：把euclidean_distances换成余弦距离或者带权重的马氏距离，就能适配文本聚类、推荐系统等场景。余弦距离在稀疏高维数据上比欧氏距离效果好很多。

约束聚类：有些业务场景要求某些样本必须分到同一簇，或者必须分开。这需要在分配步骤加约束，实现起来不复杂，但能解决实际问题。

在线聚类：数据流式到达时，可以用 Mini-Batch 思路增量更新质心。这对实时推荐、异常检测很有用。

可视化：用 PCA 或者 t-SNE 把高维数据降到二维，画散点图看聚类效果。这一步在调参时特别有用，能直观看到簇的形状和重叠情况。

我个人在实际操作中的体会是，K-Means 看起来简单，但真正用好需要理解它的假设和局限。它假设簇是球形的、大小相近的、方差相似的，一旦数据违反这些假设，结果就会很糟糕。手写一遍最大的收获不是学会了算法，而是学会了在什么时候不该用它。

最后分享一个小技巧：如果你的数据维度很高（比如上百维），先做 PCA 降维再聚类，往往效果更好、速度更快。因为高维空间里距离度量会失效（维度灾难），降维能缓解这个问题。我试过在 200 维数据上直接聚类，结果和随机差不多，降到 20 维之后簇结构就清晰了。