python实现k-means聚类算法

def initCent(dataSet , k):
 N = shape(dataSet)[1]
 cents = {}
 randIndex=[]
 #随机生成k个不重复的索引
 for i in range(k): 
 rand = random.randint(0,N)
 while rand in randIndex:
  rand = random.randint(0, N)
 randIndex.append(rand)
 #按索引取dataSet中的data作为质心
 for i in range(k):
 templist = []
 templist.append(dataSet[randIndex[i]])
 templist.append([dataSet[randIndex[i]]])
 cents[i] = templist
 return cents

#计算两个向量的欧氏距离
def calDist(X1 , X2):
 sum = 0
 for x1 , x2 in zip(X1 , X2):
 sum += (x1 - x2) ** 2
 return sum ** 0.5

#聚类
def doKmeans(dataSet , k , cents):
 #清空上一次迭代后的簇中元素，只记录质心
 for i in range(k):
 cents[i][1] = []
 for data in dataSet:
 no = 0#初始化簇标号
 minDist = sys.maxint#初始化data与k个质心的最短距离
 for i in range(k):
  dist = calDist(data , cents[i][0])
  if dist < minDist:
  minDist = dist
  no = i
 #找到距离最近的质心
 cents[no][1].append(data)
 #更新质心
 for i in range(k):
 for j in range(shape(dataSet)[0]):
  cents[i][0] = mean(cents[i][1],axis=0).tolist()
 return cents

#判断两次聚类的结果是否相同
def isEqual(old , new):
 for i in range(len(old)):
 if(old[i] != new[i][0]):
  return 0
 return 1

#主方法
def kmeans_main(dataSet,k):
 cents = initCent(dataSet, k)
 for x in range(1000):
 oldcents = []
 #拷贝上一次迭代的结果
 for i in cents.keys():
  oldcents.append(cents[i][0])
 newcents = doKmeans(dataSet, k , cents)
 #若相邻两次迭代结果相同，算法结束
 if isEqual(oldcents , newcents)>0:
  break
 cents = newcents
 return cents

def draw(cents):
 color = [ 'y', 'g', 'b']
 X = []
 Y = []
 Z = []
 fig = plt.figure()
 ax = Axes3D(fig)
 for i in cents.keys():
 X.append(cents[i][0][0])
 Y.append(cents[i][0][1])
 Z.append(cents[i][0][2])
 ax.scatter(X, Y, Z,alpha=0.4,marker='o',color='r', label=str(i))
 for i in cents.keys():
 X = []
 Y = []
 Z = []
 data = cents[i][1]
 for vec in data:
  X.append(vec[0])
  Y.append(vec[1])
  Z.append(vec[2])
 ax.scatter(X, Y, Z, alpha=0.4,marker='o', color=color[i], label=str(i),)
 plt.show()

dataSet = loadDataSet("dataSet.txt")
cents = kmeans_main(dataSet , 3)
draw(cents)