Я хочу добавить легенду к своему сюжету.У меня есть текстовые документы, я обработал их с помощью PCA, чтобы иметь возможность построить двухмерный график, но я хочу иметь легенду, объясняющую метку каждого цвета для кластеров.
Мои данные являются оригинальными спискомстрок (текстовые документы), я использовал TFIDFVectorizer, а затем PCA.Матрица, которую я получаю от применения векторизатора. Я добавил метку для каждой строки, чтобы иметь группу, к которой принадлежит этот документ.
Я могу получить график с данными 2D из PCA, и цвета правильные (кластеризация является правильной), но я просто хочу добавить легенду, гласящую: - цвет зеленый -> doctype1 - цвет красный ---> doctype2 - ....
data = vectorizer.fit_transform(documents).todense()
pca = PCA(n_components=2).fit(data)
data2D = pca.transform(data)
kmeans = KMeans(n_clusters = 4).fit(data)
clusters = kmeans.labels_.tolist()
y_means = kmeans.predict(data)
plt.scatter(data2D[:,0], data2D[:,1], c=y_means, zorder=0)
# I used n_clusters = 4 cause I know this is the optimum number of clusters
# documents is the list of strings(documents)
# I know I use the same data to predict and fit, it just to have the right colors
Спасибо