Matplotlib散点图：如何按ID设颜色并显示全部ID图例？

我尝试使用Matplotlib绘制散点图，依据ID字段为数据点设置颜色，同时期望图例能展示所有实际用到的ID及其对应颜色。以下是我的代码：

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.pyplot import figure

x = np.random.uniform(0, 100, 2000)
y = np.random.uniform(0, 50, 2000)
ID = np.random.randint(0,100,2000)

fig, ax = plt.subplots(figsize=(10, 8),dpi = 80)
scatter = ax.scatter(x,
                     y,
                    c = ID)

ax.set_xlabel('x')
ax.set_ylabel('y')
ax.legend(*scatter.legend_elements(),
          loc="center left", 
          title='ID', 
          bbox_to_anchor=(1, 0.5)
         )    
ax.ticklabel_format(useOffset=False)
ax.tick_params(axis = 'x',labelrotation = 45)

但当前生成的图例仅显示间隔分布的ID，而非数据中实际存在的全部ID，请问如何修正以显示完整的ID图例？

解决方案 #

默认的scatter.legend_elements()会自动对数值型颜色值抽样，生成间隔分布的图例项，而非全部唯一ID。以下两种方法可以解决问题：

方法一：遍历唯一ID手动添加图例项（适配ID不连续场景） #

这种方法能精准显示所有实际出现的ID，适合ID值不连续的情况：

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.cm import get_cmap

x = np.random.uniform(0, 100, 2000)
y = np.random.uniform(0, 50, 2000)
ID = np.random.randint(0,100,2000)

# 获取并排序所有唯一ID
unique_ids = np.unique(ID)
unique_ids.sort()

# 选择颜色映射，确保颜色数量匹配唯一ID数量
cmap = get_cmap('viridis', len(unique_ids))

fig, ax = plt.subplots(figsize=(10, 8), dpi=80)
handles = []

# 遍历每个ID，绘制对应数据点并收集图例项
for idx, id_val in enumerate(unique_ids):
    mask = ID == id_val
    scatter = ax.scatter(x[mask], y[mask], color=cmap(idx), label=str(id_val))
    handles.append(scatter)

ax.set_xlabel('x')
ax.set_ylabel('y')
ax.legend(handles=handles, loc="center left", title='ID', bbox_to_anchor=(1, 0.5))
ax.ticklabel_format(useOffset=False)
ax.tick_params(axis='x', labelrotation=45)

plt.show()

方法二：指定legend_elements()的num参数（适配ID连续场景） #

如果ID是连续数值范围（比如0-99），可以直接指定num参数为唯一ID的数量，强制生成全部图例项：

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

x = np.random.uniform(0, 100, 2000)
y = np.random.uniform(0, 50, 2000)
ID = np.random.randint(0,100,2000)

fig, ax = plt.subplots(figsize=(10, 8),dpi = 80)
scatter = ax.scatter(x, y, c=ID, cmap='viridis')

ax.set_xlabel('x')
ax.set_ylabel('y')
# 设置num参数为唯一ID的数量
handles, labels = scatter.legend_elements(num=np.unique(ID).size)
ax.legend(handles, labels, loc="center left", title='ID', bbox_to_anchor=(1, 0.5))

ax.ticklabel_format(useOffset=False)
ax.tick_params(axis = 'x',labelrotation = 45)

plt.show()

说明 #

• 方法一灵活性更高，无论ID是否连续都能准确显示所有实际存在的ID；
• 方法二更简洁，但需要确保ID的范围明确，避免出现多余或遗漏的图例项。

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