[笔记] 使用numpy手写k-means算法
代码包括数据生成、可视化。
注意:下面代码仅供参考,实际使用还需加上一些约束,如迭代次数需要有个最大值,等等。
import numpy as np
from matplotlib import pyplot as plt
# - generate random data
def generate_data(n_point_per_cate, center_point_list):
"""
n_point_per_cate:
point number per category
center_point_list:
center point list
"""
points_list = []
for point in center_point_list:
points_list.append(np.random.randn(n_point_per_cate, 2) + np.array(point))
return np.concatenate(points_list, axis=0)
# - generate random data
data = generate_data(100, [[3,4], [10,-4], [-5,0]])
data.shape
(300, 2)
# - visulize data
plt.scatter(data[:,0], data[:,1])
# - k-means function
def kmeans(data, K):
"""
data: input data
K: category number
"""
n,d = data.shape
cate_list = np.zeros(n)
# - random centroid
centroid_list = np.random.randn(K,d)
is_ok = False
lr = 0.5
while not is_ok:
for j in range(n):
nearest_centeroid_index = None
nearest_centeroid_distance = float('inf')
for k in range(K):
dist = np.linalg.norm(centroid_list[k] - data[j])
if dist < nearest_centeroid_distance:
nearest_centeroid_distance = dist
nearest_centeroid_index = k
cate_list[j] = nearest_centeroid_index
# - update centroid_list
last_centroid_list = centroid_list.copy()
for j in range(K):
new_centroid = np.mean(data[cate_list==j], axis=0)
centroid_list[j] = centroid_list[j]*lr + new_centroid*(1-lr)
print('centroid_list=', centroid_list)
# - visualize
plt.scatter(data[:,0], data[:,1], c=cate_list)
plt.plot(centroid_list[:,0], centroid_list[:,1], 'r+')
plt.show()
# - check if need more update
diff = np.linalg.norm(np.linalg.norm(centroid_list-last_centroid_list, axis=0))
print('diff=', diff)
if diff < 0.1:
is_ok = True
kmeans(data, K=3)
爱生活,爱拉风