Чтобы упростить процесс, я бы использовал переменную distances
из функции k_means()
, возвращая ее:
def k_means(dataset, k):
centroids = init_cent(dataset, k)
distances = calc_dist_euclidean(dataset, centroids)
cluster_assigning = np.array([dataset[distances == k].mean(axis=0) for k in range(centroids.shape[0])])
return centroids, distances, cluster_assigning
Затем вы можете построить свои центроиды и точки данных на основе соответствующего кластера, определив массив numpy colors
, содержащий каждый цвет, и создав массив значений numpy для каждой из ваших точек данных, используя colors[distances]
:
df = pd.read_csv('bristol_vacation_rentals_2016.csv')
dataset = df[['latitude', 'longitude']].values
k = 3
colors = np.array(['r', 'g', 'b'])
centroids, distances, cluster_assigning = k_means(dataset, k)
plt.subplot(121)
plt.scatter(dataset[:, 0], dataset[:, 1], c=colors[distances], s=2)
plt.scatter(centroids[:, 0], centroids[:, 1], c=colors, s=50)
plt.xlabel('latitude')
plt.ylabel('longitude')
plt.subplot(122)
plt.scatter(dataset[:, 0], dataset[:, 1], c=colors[distances], s=2)
plt.scatter(cluster_assigning[:, 0], cluster_assigning[:, 1], c=colors, s=50)
plt.xlabel('latitude')
plt.ylabel('longitude')
plt.show()