Показать легенду, что matplotlib динамически создан

Question

Мой df имеет 4 столбца: x, y, z и группировка. Я создал трехмерный график с назначенным цветом каждой точки в зависимости от того, к какой группе она принадлежит в этом ряду. Для справки, «группировка» может быть любым числом от 1 до 6. Код показан ниже:

fig = plt.figure()
ax = Axes3D(fig)
ax.scatter3D(df.x, df.y, df.z, c=df.grouping)
plt.show()

Я хотел бы показать легенду на графике, которая показывает, какой цвет принадлежит к какой группе. Ранее я использовал Seaborn для двухмерного графика, и легенда была нанесена автоматически. Как мне добавить эту функцию с помощью matplotlib?

Answer 1

Если значения для отображения цветов являются числовыми, решение может быть простым:

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

a = np.random.rand(3,40)
c = np.random.randint(1,7, size=a.shape[1])

fig = plt.figure()
ax = fig.add_subplot(111, projection="3d")
sc = ax.scatter3D(*a, c=c)
plt.legend(*sc.legend_elements())
plt.show()