Как добавить легенду к pyplot, элементы которого зависят от целочисленной переменной?

Question

Я реализую алгоритм кластеризации K-Means. Я хочу показать кластерные наборы данных (каждый разного цвета) в одной диаграмме рассеяния. Я делаю это следующим образом:

for i in range(k):
    plt.scatter(np.array(clustersets[i])[:, 0], np.array(clustersets[i])[:, 1], c=c_map(i))

, где k - количество центров (-> количество наборов кластеров) в моей диаграмме рассеяния. Теперь я хочу добавить легенду, содержащую 1 элемент для каждого из этих наборов данных. Следовательно, это зависит от заранее определенного числа k.

Как я могу добавить легенду так, чтобы она покрывала все различные элементы в моем точечном графике, не обращая внимания на то, что такое k?

Answer 1

Если я правильно понял, вы хотите, чтобы легенды шли от 0 до k.Вы можете использовать параметр label с указателем поля %d, чтобы указать набор данных.

for i in range(k):
    plt.scatter(np.array(clustersets[i])[:, 0], 
                np.array(clustersets[i])[:, 1], c=c_map(i), label='Dataset %d' %i)
plt.legend()

Answer 2

Можете ли вы попробовать с c=KMeans(n_clusters=k).fit(X).labels_astype(float)?