K-Means 算法从零实现：手写 100 行代码，可视化 5 步迭代过程
K-Means 算法从零实现：手写 100 行代码，可视化 5 步迭代过程
当数据科学家第一次接触聚类问题时，K-Means 算法往往是他们的首选武器。这个看似简单的算法背后，隐藏着优雅的数学原理和强大的实践价值。今天，我们将抛开现成的机器学习库，用不到100行Python代码实现完整的K-Means算法，并通过动态可视化揭示其迭代过程的奥秘。
1. 算法原理与核心思想
K-Means 的核心目标是将n个数据点划分到k个簇中，使得每个点到其所属簇中心的距离平方和最小。这个过程可以分解为两个关键步骤：
- 分配阶段：将每个点分配给最近的簇中心
- 更新阶段：重新计算每个簇的中心位置
数学表达式为最小化以下目标函数：
J = Σ(每个点到其簇中心的距离²)算法执行流程如下：
初始化k个簇中心 while 簇中心变化大于阈值: 将每个点分配到最近的簇中心 重新计算每个簇的中心位置注意：K-Means对初始中心点选择敏感，通常采用k-means++初始化方法减少迭代次数
2. 环境准备与数据生成
我们使用NumPy进行矩阵运算，Matplotlib进行可视化。首先创建一个包含三个明显簇的二维数据集：
import numpy as np import matplotlib.pyplot as plt # 生成模拟数据 np.random.seed(42) cluster1 = np.random.normal(loc=[0,0], scale=0.5, size=(50,2)) cluster2 = np.random.normal(loc=[3,3], scale=0.8, size=(50,2)) cluster3 = np.random.normal(loc=[-3,3], scale=0.3, size=(50,2)) data = np.vstack([cluster1, cluster2, cluster3]) # 可视化初始数据 plt.scatter(data[:,0], data[:,1], alpha=0.6) plt.title("原始数据分布") plt.show()3. 完整K-Means类实现
下面是我们实现的K-Means类，包含核心算法和可视化功能：
class KMeans: def __init__(self, n_clusters=3, max_iter=100, tol=1e-4): self.n_clusters = n_clusters self.max_iter = max_iter self.tol = tol self.history = [] # 记录每次迭代的中心点变化 def _init_centers(self, X): # 随机初始化中心点 indices = np.random.permutation(X.shape[0])[:self.n_clusters] return X[indices] def _assign_clusters(self, X, centers): # 计算每个点到各中心的距离 distances = np.linalg.norm(X[:, np.newaxis] - centers, axis=2) # 返回最近中心的索引 return np.argmin(distances, axis=1) def _update_centers(self, X, labels): centers = np.zeros((self.n_clusters, X.shape[1])) for k in range(self.n_clusters): centers[k] = np.mean(X[labels == k], axis=0) return centers def fit(self, X): self.centers = self._init_centers(X) for _ in range(self.max_iter): old_centers = self.centers.copy() self.labels = self._assign_clusters(X, self.centers) self.centers = self._update_centers(X, self.labels) self.history.append((old_centers.copy(), self.labels.copy())) # 检查收敛条件 if np.linalg.norm(self.centers - old_centers) < self.tol: break return self def plot_iterations(self, X, rows=2, cols=3): plt.figure(figsize=(15, 10)) for i, (centers, labels) in enumerate(self.history[:rows*cols]): plt.subplot(rows, cols, i+1) for k in range(self.n_clusters): cluster_data = X[labels == k] plt.scatter(cluster_data[:,0], cluster_data[:,1], alpha=0.6) plt.scatter(centers[k,0], centers[k,1], c='black', marker='x', s=100) plt.title(f'迭代 {i+1}') plt.tight_layout() plt.show()4. 算法执行与可视化分析
现在让我们使用这个类并观察其迭代过程：
# 初始化并训练模型 kmeans = KMeans(n_clusters=3) kmeans.fit(data) # 可视化迭代过程 kmeans.plot_iterations(data) # 打印最终中心点 print("最终簇中心：") print(kmeans.centers)典型的迭代过程会呈现以下演变：
- 初始随机中心：中心点可能聚集在同一区域
- 第一次分配：形成初步的簇边界
- 中心点调整：中心向数据密集区移动
- 稳定状态：中心和簇分配不再显著变化
5. 关键问题与优化技巧
在实际应用中，我们需要注意以下几个关键问题：
- 初始中心敏感度：k-means++初始化可以显著改善结果
- 肘部法则确定K值：通过不同K值的SSE曲线拐点选择最优簇数
- 数据标准化：不同量纲的特征需要归一化处理
优化后的初始化方法实现：
def kmeans_plus_plus(X, n_clusters): centers = [X[np.random.randint(X.shape[0])]] for _ in range(1, n_clusters): distances = np.array([min([np.linalg.norm(x-c)**2 for c in centers]) for x in X]) prob = distances / distances.sum() centers.append(X[np.random.choice(X.shape[0], p=prob)]) return np.array(centers)6. 数学推导与变种算法
K-Means可以看作是一种期望最大化(EM)算法：
- E步：固定中心点，优化簇分配（最小化J）
- M步：固定簇分配，优化中心点（求导令∂J/∂c=0）
常见变种包括：
| 算法变种 | 改进点 | 适用场景 |
|---|---|---|
| K-Medoids | 使用中位数代替均值 | 抗噪声数据 |
| Mini-Batch | 使用数据子集更新 | 大规模数据 |
| Fuzzy C-Means | 软分配（概率隶属度） | 重叠簇场景 |
7. 实战建议与性能考量
在实际项目中应用K-Means时：
# 最佳实践示例 from sklearn.preprocessing import StandardScaler from sklearn.metrics import silhouette_score # 数据标准化 scaler = StandardScaler() X_scaled = scaler.fit_transform(data) # 寻找最优K值 silhouette_scores = [] for k in range(2, 6): kmeans = KMeans(n_clusters=k).fit(X_scaled) score = silhouette_score(X_scaled, kmeans.labels) silhouette_scores.append(score) # 可视化K值选择 plt.plot(range(2,6), silhouette_scores, marker='o') plt.xlabel('Number of clusters') plt.ylabel('Silhouette Score') plt.show()关键性能指标对比：
- 计算复杂度：O(nkI*d)，n样本数，k簇数，I迭代次数，d特征数
- 内存需求：只需存储中心点和临时距离矩阵
- 收敛速度：通常10-20次迭代即可收敛
8. 扩展应用与创新思路
K-Means的变通应用场景：
- 图像压缩：将像素颜色聚类为k种代表色
- 异常检测：远离所有中心的点视为异常
- 特征工程：簇编号作为新特征
创新改进方向：
- 自适应K值：根据数据密度动态调整簇数
- 核方法：通过核函数处理非线性可分数据
- 流数据：增量式更新中心点
# 图像压缩示例 from sklearn.cluster import KMeans from PIL import Image def compress_image(image_path, n_colors=16): image = Image.open(image_path) pixels = np.array(image).reshape(-1, 3) kmeans = KMeans(n_clusters=n_colors).fit(pixels) compressed = kmeans.cluster_centers_[kmeans.labels_].reshape(image.size[1], image.size[0], 3) return Image.fromarray(compressed.astype('uint8'))9. 算法局限与替代方案
K-Means的固有局限：
- 假设簇是凸形且各向同性
- 对噪声和异常值敏感
- 需要预先指定簇数量
替代算法对比：
| 算法 | 优势 | 劣势 |
|---|---|---|
| DBSCAN | 自动确定簇数，处理任意形状 | 参数敏感，高维效果差 |
| 高斯混合模型 | 软聚类，概率输出 | 计算复杂，可能过拟合 |
| 谱聚类 | 能发现复杂结构 | 内存消耗大，扩展性差 |
10. 工程实践中的技巧
在真实业务场景中应用K-Means时：
- 特征选择：使用PCA降维提高聚类效果
- 评估指标：结合轮廓系数和肘部法则
- 并行化：利用Spark MLlib处理海量数据
- 增量学习：partial_fit方法处理流式数据
一个完整的聚类分析流程：
# 完整分析流程示例 def full_analysis(X, max_k=10): # 数据预处理 X = StandardScaler().fit_transform(X) X = PCA(n_components=0.95).fit_transform(X) # 寻找最优K results = [] for k in range(2, max_k+1): kmeans = KMeans(n_clusters=k).fit(X) results.append({ 'k': k, 'inertia': kmeans.inertia_, 'silhouette': silhouette_score(X, kmeans.labels_) }) # 可视化评估 fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12,4)) ax1.plot([r['k'] for r in results], [r['inertia'] for r in results]) ax1.set_title('Elbow Method') ax2.plot([r['k'] for r in results], [r['silhouette'] for r in results]) ax2.set_title('Silhouette Score') plt.show() return results11. 可视化进阶技巧
除了基本的散点图，我们还可以使用更多高级可视化：
def advanced_visualization(X, labels, centers): plt.figure(figsize=(12,5)) # 原始数据与簇分布 plt.subplot(121) for k in range(len(centers)): cluster_data = X[labels == k] plt.scatter(cluster_data[:,0], cluster_data[:,1], alpha=0.5) plt.scatter(centers[:,0], centers[:,1], c='black', marker='X', s=200) # Voronoi图展示决策边界 plt.subplot(122) x_min, x_max = X[:,0].min()-1, X[:,0].max()+1 y_min, y_max = X[:,1].min()-1, X[:,1].max()+1 xx, yy = np.meshgrid(np.linspace(x_min, x_max, 500), np.linspace(y_min, y_max, 500)) Z = np.array([np.argmin([np.linalg.norm(x-c) for c in centers]) for x in np.c_[xx.ravel(), yy.ravel()]]) Z = Z.reshape(xx.shape) plt.contourf(xx, yy, Z, alpha=0.3) plt.scatter(X[:,0], X[:,1], c=labels, alpha=0.6) plt.scatter(centers[:,0], centers[:,1], c='black', marker='X', s=200) plt.tight_layout() plt.show()12. 数学深度解析
从优化角度理解K-Means：
目标函数：
J = Σ_{i=1}^n Σ_{k=1}^K r_{ik} ||x_i - μ_k||^2其中r_{ik}是指示变量（点i是否属于簇k）
求导过程：
∂J/∂μ_k = 2Σ_{i=1}^n r_{ik}(μ_k - x_i) = 0 => μ_k = Σ r_{ik}x_i / Σ r_{ik}这正好对应我们的更新步骤：新的中心是所有属于该簇的点的平均值。
13. 多维数据与降维处理
对于高维数据，我们可以结合降维技术：
from sklearn.decomposition import PCA # 高维数据聚类示例 def high_dim_clustering(X, n_clusters=3): # 降维可视化 pca = PCA(n_components=2) X_pca = pca.fit_transform(X) # 聚类 kmeans = KMeans(n_clusters=n_clusters).fit(X) # 可视化 plt.scatter(X_pca[:,0], X_pca[:,1], c=kmeans.labels_) plt.title('PCA Projection with Cluster Labels') plt.show() return kmeans14. 时间序列与轨迹聚类
K-Means也可以应用于时间序列数据：
from scipy.spatial.distance import euclidean from fastdtw import fastdtw # 动态时间规整距离 def time_series_clustering(series_list, n_clusters=3): # 计算距离矩阵 distance_matrix = np.zeros((len(series_list), len(series_list))) for i in range(len(series_list)): for j in range(i+1, len(series_list)): distance, _ = fastdtw(series_list[i], series_list[j]) distance_matrix[i,j] = distance distance_matrix[j,i] = distance # 谱聚类 from sklearn.cluster import SpectralClustering model = SpectralClustering(n_clusters=n_clusters, affinity='precomputed') labels = model.fit_predict(distance_matrix) return labels15. 总结与未来方向
通过这次从零实现，我们深入理解了K-Means的每个细节。现代机器学习框架虽然提供了现成的实现，但亲手编写算法仍然是掌握其精髓的最佳方式。未来可以探索：
- 与深度学习结合的自适应聚类
- 处理流式数据的在线聚类算法
- 融合领域知识的半监督聚类方法
完整的实现代码和可视化工具可以轻松扩展到实际项目中，成为数据分析工具箱中的利器。