Matplotlib: 散点图中的标记和颜色图例

Question

我正在尝试为下面的数据框绘制一个散点图。我想用不同的标记（比如圆形、方形等）来表示不同的阶段，用不同的颜色（比如红色、蓝色等）来表示不同的产品。到目前为止，我已经做到了这一点，但我在显示一个图例（说明标记和颜色的部分）上遇到了困难。

这是我写的代码：

df = pd.DataFrame([[1500,24,'open','drive'],
                   [2900, 30, 'open', 'walk'],
                   [1200, 50, 'closed', 'drive'],
                   [4000, 80, 'open', 'air'],
                   [8000, 70, 'ongoing', 'air'],
                   [6100, 40, 'ongoing', 'walk'],
                   [7200, 85, 'closed', 'drive'],
                   [3300, 25, 'closed', 'drive'],
                   [5400, 45, 'open', 'walk'],
                   [5900, 53, 'open', 'air']])
df.columns = ['Cost','Duration','Stage','Product']

label_encoder = LabelEncoder()
markers = {0: 'o', 1: 's', 2: '^'}
df['Product_encoded'] = label_encoder.fit_transform(df['Product'])
df['Stage_encoded'] = label_encoder.fit_transform(df['Stage'])
df['Stage_encoded']= df['Stage_encoded'].map(markers)
colors = ('red', 'blue', 'lightgreen', 'gray', 'cyan')

X= np.array(df)
for idx,cl in enumerate(np.unique(df['Stage_encoded'])):
    plt.scatter(x=X[df['Stage_encoded']== cl,0],y=X[df['Stage_encoded']== cl,1],marker=cl,c=[colors[i] for i in X[df['Stage_encoded'] == cl, 4]])
    plt.legend()

这段代码生成了图像（见下方），并为每个点赋予了合适的颜色和标记，但我想要显示图例（标记和颜色的说明）。

我的代码输出是这样的

Answer 1

我调整了一下这个循环，让它更容易处理一些：

for stage in df["Stage"].unique(): # for every unique stage
    subD = df[df["Stage"]==stage] # get the data for that specific case
    plt.scatter(x=subD["Cost"], # cost on x axis
                y=subD["Duration"], # duration on y aaxis
                marker=subD["Stage_encoded"].unique()[0], # marker from what you defined
                # color will be automatically changing, no need to specify it in this case
                label=subD['Stage_encoded'].unique()[0]) # add a label!
plt.legend() # legend on the outside! otherwise it will turn on or off each loop!

这是图表，还有图例哦；)