plt.scatter()和plt.legend()的记录

最近在画图的时候,需要在右上角把类别和颜色对应关系画出来.因为刚开始是一个点一个点scatter的,代码如下:

data = []
lab =[]
markers = ["+", "*", "<", ".", "o", ">", ",", "p", "3"]
colors = ["red", "green", "orange", "yellow"]

for i in range(0,len(data)):
'''
遍历data中的每个点,然后将其scatter出来.
data是一个[79944,2]的矩阵
lab是一个[98844,1]的矩阵

'''
    x,y = data[i]
    path = paths[i]
    lab = int(labels[i])
    #judge_move(x, y, path)
    print("start plot {}".format(i))
    plt.scatter(x,
                y,
                label=lab,
                marker=markers[lab],
                color=colors[lab],
                # s=size,
                # linewidths=linewidths,
                # edgecolors="black"
                )

这样做耗时且不能将右上角画出类别和颜色对应关系.

改进之后:

'''
cmap='plasma'
cmap='Oranges'
cmap='summer'
cmap='Spectral'
cmap='Set1'
cmap='rainbow'
'''
classes = ["other","Normal","STR","TUM"]
end = time.time()
x = data[:,0]
y = data[:,1]

scatter = plt.scatter(x,y, c = labels, cmap ='rainbow',s=10)
plt.legend(handles = scatter.legend_elements()[0], labels=classes)