Kmeans algorithm

Kmeans.py

#coding=utf-8
import random,math,pylab

def dis(p1,p2):#compute distance between p and c
    f=lambda (x,y):(x-y)**2
    return sum(map(f,zip(p1,p2)))

def update_center(sets):#sets=[(),(),...]
    n=len(sets)
    x,y=zip(*sets)
    center=(sum(x)/n,sum(y)/n)
    return center

def cluster(points,center): #points / center =[(),(),...] 
    set_c={}
    for c in center:
        set_c[c]=[]  #initial empty setC 
    for p in points:
        minD=30000
        cen=(0,0)#initial cen to (0,0)
        for c in center:
            d=dis(p,c)
            if d<minD: # p belongs to cen
                minD=d
                cen=c
        set_c[cen].append(p) #append p to setC
    return set_c

def cost(center,new_center):
    f=lambda (x,y):x**2+y**2
    score_c=sum(map(f,center))
    score_new=sum(map(f,new_center))
    return abs(score_c-score_new)

def Initial_p(n): #initial n points
    x=[random.randint(0,100) for i in range(0,n)]
    y=[random.randint(0,100) for i in range(0,n)]
    return zip(x,y)

def Initial_c(n):
    #k=int(math.sqrt(n/2))+1
    k=int(n/2)
    x=[random.randint(0,100) for i in range(0,k)]
    y=[random.randint(0,100) for i in range(0,k)]
    return zip(x,y)

def Plot(scatter,rgb):
    return pylab.plot(zip(*scatter)[0],zip(*scatter)[1],'.',color=rgb)

if __name__=='__main__':
    #initial n-points and k-clusters k=sqrt(n/2)+1 n=30
    n=int(raw_input("input the number of initial points: "))
    points=Initial_p(n)
    center=Initial_c(n)
    initial_center=center
    
    while(1):
        clus=cluster(points,center)
        new_center=[update_center(clus[key]) for key in clus if clus[key]!=[]]#delete the too far away center
        if cost(new_center,clus.keys())<5:#center==clus.keys()
            break
        else:
            center=new_center
    center=clus.keys()
    plot1=Plot(center,'blue')
    plot2=Plot(points,'red')
    plot3=Plot(initial_center,'green')  
    pylab.show()