手写kmean

n=20
m=2
X=np.random.randn(n,m)

def kmeans(X,h,min_err=0.000001,max_iter=1000000000):
    n,m=X.shape
    err,iter=1,1
    c0=X[:h,:]
    while True:
        clusters=[np.zeros([0,m]) for i in range(h)]
        for i in range(n):
            sub=np.tile(X[i,:],[h,1])-c0
            dist=np.sum(sub**2,1)
            j=np.argmin(dist)
            clusters[j]=np.concatenate([clusters[j],X[i,None]],0)
        c1=np.zeros_like(c0)
        for i in range(h):
            c1[i,:]=np.mean(clusters[i],0)
        err=np.sum(np.sum((c0-c1)**2,1),0)
        iter+=1
        c0=c1
        print(err)
        if iter>max_iter or err<min_err:
            break
    return c0

pt=kmeans(X,10)

plt.scatter(X[:,0],X[:,1],c='r')
plt.scatter(pt[:,0],pt[:,1],c='g')