Matplotlib散点图按数组元素分类着色的高效实现问题

问题根源 #

你的待修复代码存在两个核心问题：

• data[:,2]是一维NumPy数组，直接用int(data[:,2])会报错——无法将整个数组转换为单个整数。
• 列表colors不支持数组索引，必须生成与数据点数量匹配的颜色数组，才能被plt.scatter的color参数接收。

方法一：NumPy数组索引生成颜色数组（精准匹配自定义颜色） #

这是最贴合你需求的向量化方案，直接将分类值映射到自定义颜色列表：

import numpy as np
import matplotlib.pyplot as plt

data = np.loadtxt('data.csv', delimiter = ",")
colors = ['r', 'g', 'b']

# 1. 将分类值处理为0/1/2的整数索引
category_indices = (data[:, 2] % 3).astype(int)
# 2. 用索引数组从颜色列表中生成对应颜色数组
color_array = np.array(colors)[category_indices]

fig = plt.figure()
plt.scatter(data[:,0], data[:,1], color=color_array)
plt.show()

• 步骤解析：(data[:,2]%3).astype(int)把所有分类值对3取模，转换成对应颜色列表的索引；再通过NumPy数组索引，一次性生成每个点的颜色数组，完全避免循环。

方法二：利用Matplotlib内置颜色映射（更灵活） #

如果分类数量较多，或不想手动定义颜色，可以直接用c参数传入分类数组，配合内置颜色映射：

import numpy as np
import matplotlib.pyplot as plt

data = np.loadtxt('data.csv', delimiter = ",")

fig = plt.figure()
# 直接用处理后的分类值作为颜色依据，指定颜色映射
plt.scatter(data[:,0], data[:,1], c=data[:,2]%3, cmap='tab10')
plt.colorbar(label='分类')  # 可选：添加颜色条说明分类含义
plt.show()

• 优势：无需手动维护颜色列表，Matplotlib会自动完成分类到颜色的映射，适合大规模分类场景。

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