手撕代码：用 NumPy 实现 K-means 聚类 | AI 面试题

Q: 手撕代码：用 NumPy 实现 K-means 聚类？

一、先讲清 K-means 的主循环 K-means 通过交替优化「簇分配」与「质心更新」来最小化簇内平方和（inertia）。面试里可以先说三步：每个样本分到最近质心、每个簇重新取均值、重复到质心基本不动。 二、再强调实现里的坑 代码上不要写双重 for 循环逐个算距离，可以用广播一次得到 N×K 的平方距离矩阵，复杂度是 O(NKD)。更新质心时还要显式处理空簇，否则某个簇没有样本时 mean() 会变成 NaN；初始化建议用 k-means++，让质心更分散，减少坏局部最优。 三、最后给出可运行实现 这段实现里重点看三处：距离矩阵向量化、空簇用远点替补、k-means++ 按距离平方加权采样。

核心要点

主循环三步：分配（每个样本归到最近质心）→ 更新（取簇内均值作为新质心）→ 重复直到质心稳定或达到最大迭代次数
距离用向量化广播一次算出 N×K 距离矩阵，避免双重 for 循环；比较时用平方欧氏距离即可，无需开方
空簇处理：若某簇无样本，重新挑一个离当前质心最远的点作为新质心，否则均值会得到 NaN
初始化敏感，建议用 k-means++ 按距离平方加权采样，并多次重启取惯性（inertia）最小的结果

标准回答

一、先讲清 K-means 的主循环

K-means 通过交替优化「簇分配」与「质心更新」来最小化簇内平方和（inertia）。面试里可以先说三步：每个样本分到最近质心、每个簇重新取均值、重复到质心基本不动。

二、再强调实现里的坑

代码上不要写双重 for 循环逐个算距离，可以用广播一次得到 N×K 的平方距离矩阵，复杂度是 O(NKD)。更新质心时还要显式处理空簇，否则某个簇没有样本时 mean() 会变成 NaN；初始化建议用 k-means++，让质心更分散，减少坏局部最优。

三、最后给出可运行实现

这段实现里重点看三处：距离矩阵向量化、空簇用远点替补、k-means++ 按距离平方加权采样。

python

import numpy as np

def kmeans_pp_init(X, k, rng):
    """k-means++：按距离平方加权采样初始质心，使其彼此分散。"""
    n = X.shape[0]
    centroids = [X[rng.integers(n)]]              # 随机选第一个质心
    for _ in range(1, k):
        # 每个点到已有质心的最近平方距离
        d2 = np.min(((X[:, None, :] - np.array(centroids)[None, :, :]) ** 2).sum(-1), axis=1)
        probs = d2 / d2.sum()                      # 距离越远被选中概率越大
        centroids.append(X[rng.choice(n, p=probs)])
    return np.array(centroids)

def kmeans(X, k, max_iter=100, tol=1e-4, seed=0):
    rng = np.random.default_rng(seed)
    centroids = kmeans_pp_init(X, k, rng)
    for _ in range(max_iter):
        # 1) 分配：N×K 平方欧氏距离矩阵，取最近质心
        dist2 = ((X[:, None, :] - centroids[None, :, :]) ** 2).sum(-1)
        labels = dist2.argmin(axis=1)
        # 2) 更新：逐簇取均值；空簇用离当前质心最远的点替补
        new_centroids = np.empty_like(centroids)
        for j in range(k):
            members = X[labels == j]
            if len(members) == 0:
                far = ((X - centroids[j]) ** 2).sum(-1).argmax()
                new_centroids[j] = X[far]
            else:
                new_centroids[j] = members.mean(axis=0)
        # 3) 收敛判定：质心移动小于阈值即停止
        if np.linalg.norm(new_centroids - centroids) < tol:
            centroids = new_centroids
            break
        centroids = new_centroids
    inertia = ((X - centroids[labels]) ** 2).sum()  # 簇内平方和
    return labels, centroids, inertia

if __name__ == '__main__':
    rng = np.random.default_rng(42)
    X = np.vstack([rng.normal(m, 0.5, (100, 2)) for m in ([0, 0], [5, 5], [0, 5])])
    labels, centroids, inertia = kmeans(X, k=3)
    print('inertia=', round(float(inertia), 2))
    print('centroids=\n', np.round(centroids, 2))