Python实现K-Means后，如何高效为新数据点分配聚类标签？

嘿，我懂你现在的处境——自己用Python实现了K-Means聚类，现在要给新数据贴已有的聚类标签，当前靠两层循环挨个算距离找最近质心的方法虽然能用，但总觉得不够简洁，效率也有点跟不上对吧？

先说说你现有方法的问题：两层嵌套循环的时间复杂度是O(n*k)，n是数据量，k是聚类数，当这两个数其中一个变大的时候，速度会直线下降，而且代码写起来也有点繁琐。下面给你几个更高效简洁的实现方案：

1. 用NumPy向量化操作（最推荐的纯Python高效方案） #

NumPy的底层是C实现的，向量化操作能避开Python循环的开销，速度提升非常明显，而且代码也更简洁。核心思路是用广播机制一次性计算所有数据点到所有质心的距离，再取最小距离对应的索引。

import numpy as np

def assign_cluster(cluster_dict, data):
    # 将质心和数据转换为NumPy数组
    centroids = np.array(list(cluster_dict.values()))
    data_np = np.array(data)
    
    # 计算欧氏距离的平方（开根号不影响最小值判断，省掉计算步骤更快）
    distances_sq = np.sum((data_np[:, np.newaxis] - centroids) ** 2, axis=2)
    # 找到每个数据点距离最小的质心索引，也就是聚类标签
    labels = np.argmin(distances_sq, axis=1)
    
    return labels.tolist()

这个方法不仅代码短，逻辑清晰，而且处理大规模数据时，速度比纯Python循环快几十甚至上百倍。

2. 借助Scikit-learn的现成工具（省心之选） #

如果你已经在使用机器学习生态，直接用Scikit-learn的KMeans类来做标签分配会非常省心，它内部已经做了很多优化，还能处理各种边界情况。只需要把你已经得到的质心传入模型，直接调用predict方法就行：

from sklearn.cluster import KMeans
import numpy as np

def assign_cluster(cluster_dict, data):
    centroids = np.array(list(cluster_dict.values()))
    # 初始化模型，设置质心后不需要再fit
    kmeans_model = KMeans(n_clusters=len(centroids), init=centroids, n_init=1)
    labels = kmeans_model.predict(data)
    
    return labels.tolist()

这里设置n_init=1是因为我们已经有了现成的质心，不需要重复初始化聚类了，省掉多余的计算步骤。

3. 用KD-Tree加速最近邻搜索（大数据/多聚类场景） #

如果你的数据量特别大，或者聚类数k非常多，用KD-Tree来做最近邻搜索会更高效，它能把时间复杂度降到O(n log k)，比O(n*k)的循环快很多。可以用Scipy的cKDTree实现：

from scipy.spatial import cKDTree
import numpy as np

def assign_cluster(cluster_dict, data):
    centroids = np.array(list(cluster_dict.values()))
    # 构建KD-Tree结构
    centroid_tree = cKDTree(centroids)
    # 查询每个数据点的最近邻，返回距离和对应的质心索引
    _, labels = centroid_tree.query(data, k=1)
    
    return labels.tolist()

这个方案在k很大（比如几百上千个聚类）的时候优势特别明显，能大幅减少计算时间。

总结一下 #

• 要是追求简洁高效的纯Python实现，选NumPy向量化的方法就对了；
• 要是已经在用Scikit-learn的生态，直接用它的predict方法最省心；
• 数据量极大或者聚类数超多的话，KD-Tree的方案能帮你省不少时间。

内容的提问来源于stack exchange，提问作者efsee