python - matplotlib 从 pandas 数据框中散布颜色

Question

我想做的是创建一个如下所示的 matlab 散点图，但根据 C 列中的字符串为点着色，例如，yes = 蓝色，no = 红色。

 import matplotlib.pyplot as plt
 A = pd.DataFrame.from_items([('A', [1, 2, 3]), ('B', [4, 5, 6]), ('C', ['yes', 'yes', 'no'])])


plt.scatter(A['A'], A['B'])

Answer 1

您可以创建一个(值:颜色)对的字典，并使用列表理解将这些值用作 plt.scatter 上的 color 参数。

colors = {'yes': 'b', 'no': 'r'}

plt.scatter(A['A'], A['B'], color=[colors[r] for r in A['C']])

<小时/>

编辑:添加图例的首选方法是迭代数据集中的唯一类并使用 legend 参数。因此，我们可以迭代按 C 分组的数据帧，并使用标签参数单独绘制它们。最后调用 plt.legend() 来显示它。

colors = {'yes': 'b', 'no': 'r'}

for idx, row in A.groupby('C'):
    plt.scatter(row['A'], row['B'], 
                color=[colors[r] for r in row['C']], 
                label=idx)

plt.legend()