Свяжите цвет с числом в диаграмме рассеяния matplotlib

Question

Я использую sklearn для создания кластера kmeans на основе некоторых розничных данных.

Мы используем этот кластер за кулисами для сегментирования клиентов (например, синие клиенты - отличные клиенты, зеленые клиенты имеют такую-то потребность и т. Д. c.).

Показанные точки данных имеют другой цвет в зависимости от того, в какой из 4 кластеров помещен покупатель. Но я не могу найти способ напрямую определить, какой цвет соответствует номеру сегмента (или как заставить определенные номера сегментов быть определенного цвета).

c=y в разбросе - это то место, где используется значение y (т. е. прогнозируемый сегмент для наблюдения), чтобы выбрать цвет. Всего 4 сегмента. Я просто не знаю, какой из этих четырех цветов сопоставлен с каким цветом! 1011 *

Answer 1

Вам нужно вернуть дескриптор из цветовой панели разброса и построения,

cm = ax.scatter(X[:, 0], X[:, 1], X[:, 2], c=y)
plt.colorbar(cm)

В качестве минимального примера

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

X = np.random.randn(100,3)
y = np.random.randn(100)

fig = plt.figure()
ax = Axes3D(fig)
ax.view_init(30)
cm = ax.scatter(X[:, 0], X[:, 1], X[:, 2], c=y)
plt.colorbar(cm)
plt.show()