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

你可能感兴趣的:(秋招面试,Python)