python简易实现k-means

 用dist存放所有数据到中心的距离,有n行（n组数据），k+1列（前k列分别存放到第i个类中心的距离，最后一列存放分到了第几类）

 

#!/usr/bin/env python 
# -*- coding:utf-8 -*-
import numpy as np
n = 100
x = np.arange(100)
y = np.arange(200, 300, 1)

# 1、选中心,此时假设分为两个类
k = 2
center0 = np.array([x[0],y[0]])
center1 = np.array([x[1],y[1]])


dist = np.zeros([n, k+1])
while True:
    # 2、计算距离
    for i in range(n):
        dist[i, 0] = np.sqrt((x[i]-center0[0])**2 + (y[i]-center0[1])**2)
        dist[i, 1] = np.sqrt((x[i]-center1[0])**2 + (y[i]-center1[1])**2)
        if dist[i, 0] <= dist[i, 1]: # 3、根据距离值的大小来分类
            dist[i, 2] = 0
        else:
            dist[i, 2] = 1
    # 4、 计算新的类中心
    index0 = dist[:,2] == 0 # 所有行的第三列为0
    index1 = dist[:,2] == 1 # 所有行的第三列为1
    center0_new = np.array([x[index0].mean(), y[index0].mean()]) # 逻辑值索引
    center1_new = np.array([x[index1].mean(), y[index1].mean()])
    # 5、判定结束算法
    if (center0 == center0_new).all() and (center1 == center1_new).all() :
            break
    else:
        center0 = center0_new
        center1 = center1_new
print(dist)
print(center0,center1)