如何为Matplotlib散点图生成两个图例并优化连线代码？

一、图例错误原因及修正方案 #

原代码的图例错误在于仅通过labels参数指定文本，matplotlib会自动选取绘图元素生成图例，导致混入连线元素或无法正确匹配散点的颜色/大小。正确做法是显式收集散点对象（handles），再基于这些对象创建对应图例。

修正逻辑： #

1. 绘制训练集散点时，将每个散点对象存入列表；
2. 绘制p、p1散点时，保存对应的散点对象；
3. 分别基于两组散点对象和标签创建图例，并将第一个图例添加回坐标轴（避免被第二个图例覆盖）。

二、连线代码优化 #

原代码硬编码每条连线，扩展性差。通过遍历X_new中的所有点，批量绘制到p和p1的连线，减少重复代码，同时支持后续增加更多训练点。

完整修正代码 #

import numpy as np
import matplotlib.pyplot as plt

# 生成随机数据
X_new = np.random.randint(1,20,(3,2))
p = np.random.randint(1,20,(1,2))
p1 = np.random.randint(1,20,(1,2))

color_map = {0: 'blue', 1:'green', 2: 'darkred', 3: 'black', 4:'red'}
legend1_label = {0: 'trn1', 1: 'trn2', 2: 'trn3'}
legend2_label = {0: 'p', 1: 'p1'}

# 绘制训练集散点并收集对象
trn_scatters = []
for idx, cl in enumerate(legend1_label):
    scatter = plt.scatter(x=X_new[idx, 0], y=X_new[idx, 1], 
                          c=color_map[cl], marker='.', s=100)
    trn_scatters.append(scatter)

# 批量绘制训练点到p的连线
for point in X_new:
    plt.plot([point[0], p[0,0]], [point[1], p[0,1]], 
             color='k', linestyle='--', linewidth=1)

# 批量绘制训练点到p1的连线
for point in X_new:
    plt.plot([point[0], p1[0,0]], [point[1], p1[0,1]], 
             color='r', linestyle='--', linewidth=1)

# 绘制p、p1散点并收集对象
p_scatter = plt.scatter(x=p[0,0], y=p[0,1], 
                        c=color_map[4], marker='.', s=200)
p1_scatter = plt.scatter(x=p1[0,0], y=p1[0,1], 
                         c=color_map[3], marker='.', s=200)
pp_scatters = [p_scatter, p1_scatter]

# 创建并添加双图例
legend1 = plt.legend(handles=trn_scatters, labels=list(legend1_label.values()), 
                     loc=4, title="legend1")
plt.gca().add_artist(legend1)  # 保留第一个图例
plt.legend(handles=pp_scatters, labels=list(legend2_label.values()), 
           title="legend2")

plt.show()

关键说明 #

• 图例修正：使用handles参数指定图例对应的散点对象，确保图例仅显示散点，且颜色、大小与实际绘图一致；
• 连线优化：通过循环遍历训练点，批量生成连线，代码更简洁且易于维护；
• 颜色一致性：p、p1的颜色直接复用color_map中的定义，避免硬编码导致的颜色不一致问题。

内容的提问来源于stack exchange，提问作者Leshui