python 手写kmeans聚类算法

看到一些面经中出现了手写K-means，因此自己写了一版，如有错误，欢迎指正。

输入：
raw_data: list()（例如，多个坐标点[[0, 0], [1,1],…]）
k： int 簇的个数
mse_limit: float 若更新中心点后的mse和更新前的mse，误差在该值以内，则停止迭代
early_stopping：最大迭代次数

输出：
聚类后的结果：dict(),（例如： {0: [[-0.2, -0.2], [0.1, 0.3], [0.2, 0.2], [0.2, 0]], 1: [[-0.1, 1]],…}）

import random
import numpy as np

# 初始化簇心
def get_init_centers(raw_data, k):
    return random.sample(raw_data, k)

# 计算距离
def cal_distance(x, y):
    return np.linalg.norm(np.array(x) - np.array(y))

# 将各点分配到最近的点, 并计算MSE
def get_cluster_with_mse(raw_data, centers):
    distance_sum = 0.0
    cluster = {}
    for item in raw_data:
        flag = -1
        min_dis = float('inf')
        for i, center_point in enumerate(centers):
            dis = cal_distance(item, center_point)
            if dis < min_dis:
                flag = i
                min_dis = dis
        if flag not in cluster:
            cluster[flag] = []
        cluster[flag].append(item)
        distance_sum += min_dis**2
    return cluster, distance_sum/(len(raw_data)-len(centers))

# 计算各簇的中心点，获取新簇心
def get_new_centers(cluster):
    center_points = []
    for key in cluster.keys():
        center_points.append(np.mean(cluster[key], axis=0)) # axis=0，计算每个维度的平均值
    return center_points

# K means主方法
def k_means(raw_data, k, mse_limit, early_stopping):
    old_centers = get_init_centers(raw_data, k)
    old_cluster, old_mse = get_cluster_with_mse(raw_data, old_centers)
    new_mse = 0
    count = 0
    while np.abs(old_mse - new_mse) > mse_limit and count < early_stopping : 
        old_mse = new_mse
        new_center = get_new_centers(old_cluster)
        print(new_center)
        new_cluster, new_mse = get_cluster_with_mse(raw_data, new_center)  
        count += 1
        print('mse:',np.abs(new_mse), 'Update times:',count)
    return new_cluster