如何在matplotlib的ax.scatter中设置两类颜色规则并正确显示图例

问题原因 #

你当前的图例是基于第一次批量绘制的所有散点对象points生成的，这部分散点使用Dark2配色，标签2对应的默认颜色为灰色。后续叠加绘制的黑色标签2散点没有参与图例项的生成，因此出现图例颜色和实际散点颜色不一致的问题。

解决方法 #

方案1：最小改动现有代码 #

只需要在生成图例句柄后，手动修改标签2对应的句柄颜色即可，修改后的完整代码如下：

import pandas as pd
import matplotlib.pyplot as plt
from sklearn import datasets
from collections import Counter
iris = datasets.load_iris()
X = iris.data
y = iris.target
df = pd.DataFrame(X, columns = iris.feature_names)
fig, ax = plt.subplots(figsize=(12,8))
points = ax.scatter(df.values[:,0], 
    df.values[:,1],
    c = y,
    cmap='Dark2')  

for i in range (len(y)):
    if y[i] == 2:
        ax.scatter(df.values[i,0],df.values[i,1], c = 'k') 

handles, _ = points.legend_elements()
# 新增代码：修改标签2对应的图例标记颜色为黑色
handles[2].set_color('k')
labels =sorted([f'{item}: {count}' for item, count in Counter(y).items()])
ax.legend(handles, labels, loc = "lower right",title = 'clusters')  

plt.show()

方案2：更优雅的实现（无需叠加绘制） #

直接自定义颜色映射，一次完成所有散点的绘制，避免循环渲染和重叠绘制的性能问题，代码如下：

import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.colors import ListedColormap
from sklearn import datasets
from collections import Counter
iris = datasets.load_iris()
X = iris.data
y = iris.target
df = pd.DataFrame(X, columns = iris.feature_names)
fig, ax = plt.subplots(figsize=(12,8))

# 构造自定义配色：取Dark2的前3个颜色，将标签2对应的颜色替换为黑色
base_cmap = plt.cm.get_cmap('Dark2', 3)
custom_color_list = list(base_cmap.colors)
custom_color_list[2] = 'black'
custom_cmap = ListedColormap(custom_color_list)

# 单次绘制所有散点
points = ax.scatter(
    df.values[:,0], 
    df.values[:,1],
    c = y,
    cmap = custom_cmap
)  

handles, _ = points.legend_elements()
labels = sorted([f'{item}: {count}' for item, count in Counter(y).items()])
ax.legend(handles, labels, loc = "lower right",title = 'clusters')  

plt.show()

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