如何为pandas DataFrame绘制的分类散点图添加matplotlib多色图例？

你当前的写法是将所有散点一次性传入ax.scatter()，matplotlib会将整批数据识别为单一绘图序列，因此调用plt.legend()时默认只会生成一个图例条目。

推荐按类别分批绘制散点，每个类别单独指定label参数，实现逻辑最简洁：

import pandas as pd
import matplotlib.pyplot as plt

dataset = pd.read_csv('iris.csv', header=None, names=['sepal length', 'sepal width', 'petal length', 'petal width', 'class'])

dataset.drop(index=dataset.index[dataset['class'] == 'Iris-setosa'], inplace=True)
dataset.drop(['petal width', 'sepal width'], axis='columns', inplace=True)

fig, ax = plt.subplots()
colors = {'Iris-versicolor': 'blue', 'Iris-virginica': 'red'}
labels = {'Iris-versicolor': 'aaa', 'Iris-virginica': 'bbb'}

# 按类别循环绘制散点
for cls, group in dataset.groupby('class'):
    ax.scatter(group['petal length'], group['sepal length'], 
               c=colors[cls], label=labels[cls])

# 直接调用legend即可自动读取各系列的label
plt.legend()
plt.show()

如果需要保留原有一次性绘制所有散点的逻辑，也可以手动构造图例句柄：

import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.patches import Patch

dataset = pd.read_csv('iris.csv', header=None, names=['sepal length', 'sepal width', 'petal length', 'petal width', 'class'])

dataset.drop(index=dataset.index[dataset['class'] == 'Iris-setosa'], inplace=True)
dataset.drop(['petal width', 'sepal width'], axis='columns', inplace=True)

fig, ax = plt.subplots()
colors = {'Iris-versicolor': 'blue', 'Iris-virginica': 'red'}

ax.scatter(dataset['petal length'], dataset['sepal length'], c=dataset['class'].map(colors))

# 手动构造图例元素
legend_elements = [Patch(facecolor='blue', label='aaa'),
                   Patch(facecolor='red', label='bbb')]
plt.legend(handles=legend_elements)

plt.show()

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