手写 KMeans算法并绘制动图

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


def Kmeans(center):  #【【】，【】，【】】
    #3. 第一聚类：
    for i in range(3):
        #计算每个点与三个类中心的距离：
        distance = np.sqrt(((x_train - center[i, :])**2).sum(axis=1))
        columns_name = '类{}'.format(i)
        data[columns_name] = distance

    #4. 完成第一次聚类结果：
    data['custer'] = data.loc[:, ['类0', '类1', '类2']].idxmin(axis=1)

    #5.更新类中心点：
    new_center = data.loc[:, ['平均消费周期（天）', '平均每次消费金额', 'custer']].groupby(by='custer').mean().values
    return  new_center

#绘图
def plot_image(times, data):
    plt.figure()
    for i in ['类0', '类1', '类2']:
        x = data.loc[data['custer']==i, '平均消费周期（天）']
        y = data.loc[data['custer']==i, '平均每次消费金额']
        plt.scatter(x, y)

    image_name = '第{}次聚类结果.png'.format(times)
    plt.savefig(image_name)
    plt.close('all')
    return image_name



if __name__ == '__main__':
    data = pd.read_csv('company.csv', encoding='gbk')
    #0. 生成图片的list:
    image_list = []


    # 1. 挑选出特征值：
    x_train = data.loc[:, ['平均消费周期（天）', '平均每次消费金额']]

    # 2.给出初始聚类中心： k=3
    center = np.array([[10, 100], [20, 200], [30, 300]])

    #3.完成第一次聚类
    new_center = Kmeans(center)
    print(data)
    times = 1

    image_name=plot_image(times, data)
    image_list.append(image_name)

    #第一次跌代
    while True:
        if (center==new_center).all():
            print(times)
            break
        times+=1
        center = new_center.copy()
        new_center = Kmeans(center)
        image_name = plot_image(times, data)
        image_list.append(image_name)



#读取图片
frame_list = []
for image_name in image_list:
    #读取图片：
    im = imageio.imread(image_name)

    #存放所有读进来的图：
    frame_list.append(im)


#生产gif图片
# duration: 两张图片间的间隔，建议写0.2 - 0.5,
duration= 0.7
imageio.mimsave('聚类结果.gif', frame_list, 'GIF', duration=round(duration, 2))