算法工程师面试题之手撕K-means（python）

算法思想

1. 选择K个点作为初始质心
2. repeat
   1. 计算每个样本点到每个质心的距离
   2. 将每个样本点指派到最近的质心中，形成K个簇
   3. 重新计算每个簇的质心
3. 直到簇不发生变化（分类情况不再变化）或者达到最大迭代次数

问题拆解

输入数据：训练样本集(n个样本，p个特征） train_data (n,p)，初始 k个质心 centers  (k,p)

输出数据：每个样本的 标签（n,)  以及最终的 k个质心 centers (k,p)

循环部分模块：

1. 计算样本到质心的距离，直接算整体样本。输入数据需要 train_data  和 centers。n个样本k个质心，每个样本需要分别计算与k个簇的距离，所以输出是 (n,k)
2. 将样本点指派到最近的质心中，并更新质心。返回样本点新类别 (n,), 和新的k个质心 centers
3. 判断簇是否发生变化（两种方法：①判断点的类别是否有变化 ②判断质心是否有变化），有变化就循环迭代直到没有变化

手撕代码

导入包

import numpy as np
import pandas as pd
import math

模块1：计算所有样本到质心的距离

def all_distances(train_data, centers):
    all_dis = []   # 存储所有样本到每个簇中心的距离,使用欧式距离
    for one_data in train_data:
        # 计算每个样本到k个簇的距离，所有有k个值
        one_dis = []
        for center in centers:
            # 样本和簇质心可以看成两个点(不一定是二维的哦)，求两个点的欧式距离
            one_dis.append(math.sqrt(np.sum((one_data - center) ** 2))))
        all_dis.append(one_dis)
    return all_dis

模块2：将样本点指派到距离最近的质心中，全部指派完样本后更新质心。

def update(train_data, centers):
    k = len(centers)
    p = len(train_data[0])  # p为特征维度

    new_class = []
    new_centers = []
    clusters = [[] for _ in range(k)]  # 将所有样本分门别类存储，全部存储后对每个簇求平均作为该簇的簇中心
    
    # 更新类别
    all_dis = all_distances(train_data, centers)
    for i in range(len(all_dis)):  # i相当于样本的编号
        idx = all_dis[i].index(min(all_dis[i])))
        new_class.append(idx)
        clusters[idx].append(i)  # 存样本编号，导致后根据编号到train_data中找
    
    # 更新簇中心
    for j in range(k):
        # 如果一个簇没有样本，则保持该簇中心（不是去除该簇哦）
        if len(clusters[j]) == 0:
            new_centers.append(centers[j])
        else:
            sums = np.zero((p,))  
            for i in clusters[j]:
                sums += train_data[i]
            means = sums/len(clusters[j])
            new_centers.append(means)

    return new_class, new_centers

 模块3：判断簇是否发生变化，因为上面已经计算了新的簇质心，所以直接比较簇质心。(也可以先不计算簇质心，只计算类别，判断类别是否发生变化，发生变化再计算簇中心)。

# 比较两个数组是否相同，一个元素一个元素比较
def if_changed(new_centers, centers):
    k, p = len(centers), len(centers[0])
    for i in range(k):
        for j in range(p):
            if new_centers[i][j] != centers[i][j]:
                return False
    return True

 K-Means 主函数

def k_means(train_data, centers):
    changed = True
    while changed:
        new_class, new_centers = update(train_data, centers)
        changed = if_changed(new_centers, centers)
    return new_class, new_centers

测试数据运行

if __name__ == "__main__":
    data = pd.DataFrame([["张三",12,15,13,28,24],["李四",7,11,10,19,21],["王五",12,14,11,27,23],["赵六",6,7,4,13,20],["刘七",13,14,13,27,25]])
    # 分成三个簇，初始簇表示为前三个
    train_data = data.iloc[:,1:].values
    centers = data.iloc[:3,1:].values
    new_class, new_centers = k_means(train_data, centers)
    print("分类结果:", new_class)
    print("类别：", new_class)
    print("簇中心：")
    for i in range(len(new_centers)):
        print(new_centers[i])

输出结果