import numpy as np
from sklearn.cluster import KMeans
import matplotlib.pyplot as plt
X = np.array([[20, -40], [300, -320], [300, -380], [360, -100],
[360, -320], [80, -320], [300, -100], [80, -40],
[360, -380], [80, -100], [20, -380], [20, -100],
[20, -320], [360, -40], [300, -40], [80, -380]])
def get_points_by_cl_num(clusters, num):
return X[np.where(clusters == num)]
n = 4
kmeans = KMeans(n_clusters=n)
kmeans.fit(X)
from random import choice
colors = ['red', 'green', 'blue', 'yellow']
for cl in range(n):
print('Points in cluster {}:'.format(cl))
print(get_points_by_cl_num(kmeans.labels_, cl), '\n')
for x, y in get_points_by_cl_num(kmeans.labels_, cl):
plt.scatter(x, y, c=colors[cl])
plt.show()