matplotlib绘图中使用多种颜色

Question

我有一个二维数据点（x，y），它被分为三类（0,1,2）numpy阵列。

a = array([[ 1, 2, 3, 4, 5, 6, 7, 8 ], 
      [ 9, 8, 7, 6, 5, 4, 3, 2 ]]) 

class = array([0, 2, 1, 1, 1, 2, 0, 0]) 

我的问题是如果我可以用多种颜色绘制这些点。我想这样做：

colors = list() 
for i in class: 
    if i == 0: 
     colors.append('r') 
    elif i == 1: 
     colors.append('g') 
    else: 
     colors.append('b') 

print colors 
['r', 'b', 'g', 'g', 'g', 'b', 'r', 'r'] 

pp.plot(a[0], a[1], color = colors) 

Answer 1

我假设你想绘制不同的点。在这种情况下， 如果定义了numpy的数组：

colormap = np.array(['r', 'g', 'b']) 

则可通过colormap[categories]产生的颜色的数组：

In [18]: colormap[categories] 
Out[18]: 
array(['r', 'b', 'g', 'g', 'g', 'b', 'r', 'r'], 
     dtype='|S1') 

---

import matplotlib.pyplot as plt 
import numpy as np 

a = np.array([[ 1, 2, 3, 4, 5, 6, 7, 8 ], 
       [ 9, 8, 7, 6, 5, 4, 3, 2 ]]) 

categories = np.array([0, 2, 1, 1, 1, 2, 0, 0]) 

colormap = np.array(['r', 'g', 'b']) 

plt.scatter(a[0], a[1], s=50, c=colormap[categories]) 
plt.show() 

产量