python - 显示 matplotlib 动态创建的图例

Question

我的 df 有 4 列:x、y、z 和分​​组。我创建了一个 3D 图，每个点的指定颜色由它在该行中所属的分组决定。作为引用，“分组”可以是 1 到 6 之间的任意数字。代码如下所示:

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

我想在绘图上显示一个图例，显示哪种颜色属于哪个分组。之前，我使用 Seaborn 绘制 2D 绘图，并自动绘制图例。如何使用 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()