如何为Matplotlib基于0/1值的双色散点图设置自定义图例？

我明白你遇到的问题啦——直接用c=z加label=z时，matplotlib没法自动识别出这是两个不同的类别，所以只会显示一个图例条目。下面给你两种简单的解决办法，既能分开显示0和1的图例，还能自定义标签和颜色：

方法一：分开绘制不同类别的点 #

最直观的方式就是把z=0和z=1的点分别绘制，各自指定标签和颜色。这样matplotlib会自动为每个scatter生成单独的图例项：

import matplotlib.pyplot as plt
import numpy as np

plt.figure(figsize=(15,5))
x = np.linspace(0, 12, 12)
y = np.random.random(12)
z = [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]

# 生成两类点的掩码
mask_zero = np.array(z) == 0
mask_one = np.array(z) == 1

# 分别绘制，自定义颜色和标签
plt.scatter(x[mask_zero], y[mask_zero], c='purple', label='zero')
plt.scatter(x[mask_one], y[mask_one], c='yellow', label='one')

plt.legend()
plt.show()

方法二：用自定义颜色映射+手动构建图例 #

如果你不想拆分数据，也可以保留一次绘制所有点的方式，然后手动创建图例元素。这种方式适合类别较多的场景：

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import Patch

plt.figure(figsize=(15,5))
x = np.linspace(0, 12, 12)
y = np.random.random(12)
z = [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]

# 自定义颜色映射，对应0和1的颜色
custom_cmap = plt.cm.colors.ListedColormap(['purple', 'yellow'])
plt.scatter(x, y, c=z, cmap=custom_cmap)

# 手动创建图例的色块元素
legend_items = [
    Patch(facecolor='purple', label='zero'),
    Patch(facecolor='yellow', label='one')
]
plt.legend(handles=legend_items)

plt.show()

为什么原来的代码不生效？ #

当你把整个z列表传给label参数时，matplotlib会把它当成一个整体的标签，而不会自动解析里面的不同类别值。所以必须通过拆分绘制或者手动构建图例的方式，告诉matplotlib这是两个独立的类别。

内容的提问来源于stack exchange，提问作者Rieneke van Noort