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 是个 array，里面只有0，1两种取值
# X_train 也是 array 有两列，X_train[i,0] 代表第 i 行，第 0 列, X_train[i][0] 也一样
# scatter 参数中，lw:点宽，s:不知道，color:'r'是红色, marker:'^'代表三角形

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',marker='^')


plt.xlabel('Age')

plt.ylabel('Estimated Salary')

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

# ('0','1') 代表与 (s1, s2) 对应的 label
# 也可写成 ('Purchased', 'UnPurchased')
plt.legend((s1,s2),('0','1') ,loc = 'best')