python手写kmeans

kmeans是数据挖掘中最简单的一个算法， 正是因为它简单，在很多面试中有的面试官会要求你手写一段kmeans的代码。笔者结合自己对kmeans的理解，采用西瓜书中数据手写了一个kmeans算法，跟网上版本略有出入，但思路是差不多的，可以供大家思考。
首先来说一下kmeans的原理，其实就是根据距离将样本划分为k类，最终实现类内距离最短和类间距离最大，其他原理比较简单就不多说了，大家如果有疑问自己百度一下相关内容。下面贴出代码。

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

#知识点用map函数将str转化为floatlist
base_data = []
with open('xigua.txt') as f:
    for line in f.readlines()[1:]:
        data_line = line.strip().split(',')
        data_line = list(map(float, data_line))
        base_data.append(data_line)
print(base_data)


def cal_distance(x, y):
    x = np.array(x)
    y = np.array(y)
    return np.sqrt(np.sum(np.power(x-y, 2)))
# distance = caldistance(base_data[0], base_data[1])
# print(distance)


def cal_eplision(x, y, eplision):
    if x - y <= eplision:
        return 1
    else:
        return 0



def set_init_center(data_set, k):
    data_set = np.array(data_set)
    init_center = np.zeros((k, data_set.shape[1]))
    max_feature = np.max(data_set, axis=0)
    min_feature = np.min(data_set, axis=0)
    for i in range(k):
        for j in range(data_set.shape[1]):
            init_center[i, j] = min_feature[j] + np.random.random() * (max_feature[j] - min_feature[j])
    #array可以直接整行或者整列运行
    return init_center
# init_center = set_init_center(base_data, 5)
# print(init_center)


def kmeans(data_set, k):
    assert k <= len(data_set)
    data_set = np.array(data_set)
    data_with_label = np.zeros(len(data_set))
    init_center = set_init_center(data_set, k)
    flag = 1
    optimized_distance = 0
    while flag:
        sum_point = np.zeros((1, k))
        sum_feature = np.zeros((k, len(data_set[0])))
        sum_distance = np.zeros((1, k))
        for i in range(len(data_set)):
            min_center = 0
            min_distance = 10e+6
            for j in range(len(init_center)):
                distance_tmp = cal_distance(init_center[j], data_set[i])
                if distance_tmp <= min_distance:
                    min_center = j
                    min_distance = distance_tmp
            data_with_label[i] = min_center
            sum_point[0, min_center] += 1
            sum_feature[min_center, :] += data_set[i, :]
            sum_distance[0, min_center] += min_distance
        sum_point = np.repeat(sum_point, len(data_set[0]), axis=0).reshape(k, len(data_set[0]))
        new_center = sum_feature/sum_point
        # np.mean() array也可以用逻辑切片
        # print(sum_distance)
        sum_distance_tmp = sum(sum_distance[0])
        # print(sum_distance_tmp)
        if not cal_eplision(sum_distance_tmp, optimized_distance, 10e-6):
            flag = 0
        optimized_distance = sum_distance_tmp
    return new_center, data_with_label


def plot_kmeans(base_data, k):
    new_center, data_with_label = kmeans(base_data, k)
    base_data = np.array(base_data)
    new_base_data = np.insert(base_data, 0, values=data_with_label, axis=1)
    new_base_data = pd.DataFrame(new_base_data)
    fig, ax = plt.subplots(figsize=(10, 8))
    corlor_set = ['r', 'b', 'g', 'm', 'y', 'k']
    for i in range(k):
        new_base_data_tmp = new_base_data[new_base_data[0] == i]
        ax.scatter(new_base_data_tmp[1], new_base_data_tmp[2], c=corlor_set[i])
    plt.show()


plot_kmeans(base_data, 3)