给散点图加legend图例

首先从文件导入数据，并进行格式化为numpy矩阵

import numpy as npdef file2array(filename):file = open(filename)file_lines = file.readlines()numberOfLines = len(file_lines)dataArray = np.zeros((numberOfLines, 3))labels = []index = 0for line in file_lines:line = line.strip() # 参数为空时，默认删除开头、结尾处空白符（包括'\n', '\r',  '\t',  ' ')formLine = line.split('\t')dataArray[index,:] = formLine[0:3]labels.append((formLine[-1]))index += 1label_id = []dic = {}index = 1for i in range(len(labels)):x = dic.get(labels[i], -1)label_id.append(x)if (label_id[i] == -1):label_id[i] = indexdic[labels[i]] = indexindex += 1return dataArray, label_iddata, labels = file2array('datingTestSet.txt')

然后开始对数据进行分标签绘图，开始是这样的

%matplotlib inline
import matplotlib.pyplot as plt
plt.scatter(data[:,0], data[:,1], c=15*np.array(labels))

然后给他添加legend
思路是把三种不同标签的图分开，分成三个子图，画在一张图里面

def scatterLegend(data, labels, x, y):type1 = []type2 = []type3 = []for i in range(len(labels)):if labels[i] == 1:type1.append(np.array(data[i]))elif labels[i] == 2:type2.append(np.array(data[i]))else:type3.append(np.array(data[i]))type1 = np.array(type1)type2 = np.array(type2)type3 = np.array(type3)g1 = plt.scatter(type1[:, x], type1[:, y], c='red')g2 = plt.scatter(type2[:, x], type2[:, y], c='yellow')g3 = plt.scatter(type3[:, x], type3[:, y], c='blue')plt.legend(handles=[g1, g2, g3], labels=['not at all', 'a small doses', 'a large doses'])scatterLegend(data, labels, 0, 1)

结果变成了这样