python在scatter中添加图例_plt.scatter 二分类数据 legend 添加图例

画图代码参考的帖子：~~~~~~点这~~~~~~

一个二分类任务，根据年龄和薪资估计其是否会买车，

数据集有两个特征，Age 和 Salary

标签是 Purchased ，两个取值(0或1)

横轴 Age, 纵轴 Salary，把这些离散的点画出来，并根据标签取值上色

环境 jupyter, python 3.7

# This tells Jupyter to set up Matplotlib so it uses Jupyter’s own backend.

%matplotlib inline

import matplotlib.pyplot as plt

# 根据训练集标签 y 的取值确定所画点的颜色

# y_train 是个 list，里面只有0，1两种取值

# X_train 有两列,X_train[i,0] 代表第 i 行，第 0 列

# scatter 参数中，lw:点宽，s:不知道，color:颜色，'r'是红色

for i in range(len(y_train)):

if y_train[i] == 0:

s1 = plt.scatter(X_train[i,0], X_train[i,1], s=50, lw=3, color = 'r')

elif y_train[i] == 1:

s2 = plt.scatter(X_train[i,0], X_train[i,1], s=50, lw=3, color = 'g')

plt.xlabel('Age')

plt.ylabel('Estimated Salary')

plt.title(" Logistic Regression (Training Set)")

plt.legend((s1,s2),('0','1') ,loc = 'best')