使用自己定义的色彩方案 ListedColormap 作 plt.scatter() 的参数

参考链接： https://blog.youkuaiyun.com/MaxxiChen/article/details/106137315.
这里以鸢尾花数据集为例，

import pandas as pd
import numpy as np
import tensorflow as tf
import matplotlib as mpl
import matplotlib.pyplot as plt

plt.rcParams['font.sans-serif'] = "SimHei"
plt.rcParams['axes.unicode_minus'] = False

# 第一步：加载数据
TRAIN_URL = "http://download.tensorflow.org/data/iris_training.csv"
train_path = tf.keras.utils.get_file(TRAIN_URL.split('/')[-1], TRAIN_URL)
df_iris = pd.read_csv(train_path, header=0)  # 表示第一行数据作为列标题

iris = np.array(df_iris)  # 将二维数据表转换为 Numpy 数组, (120, 5), iris的训练集中有120条样本,

# 提取属性和标签
train_x = iris[:, 0:2]  # 我们只取出前两列属性, (120, 2), 花萼的长度和花萼的宽度
train_y = iris[:, 4]  # 取出最后一列作为标签值, (120,)
# 2.4 可视化样本
cm_pt = mpl.colors.ListedColormap(["blue", "red", "green"])
# 取出花萼长度和花萼宽度作为样本点的横坐标和纵坐标, 根据样本点的标签值确定样本的颜色, 设置色彩方案为cm_pt
plt.scatter(train_x[:, 0], train_x[:, 1], c=train_y, cmap=cm_pt)
# 之前的例程中, 我们使用的都是 matplotlib 中预设的色彩方案,这里我们使用自己定义的色彩方案,
# 在散点图中, 蓝色的是山鸢尾, 红色的是变色鸢尾, 绿色的是维吉尼亚鸢尾.
plt.show()

上述代码中，train_y是样本标签值，有三种取值
样本值 0 — blue，
样本值 1 — red，
样本值 2 — green

运行代码，结果如下：