Używanie wielu kolorów w macierzy matplotlib

Question

Mam szeregową tablicę punktów danych 2D (x, y), które są podzielone na trzy kategorie (0,1,2).

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]) 

Moje pytanie brzmi, czy mogę wydrukować te punkty z wieloma kolorami. Chciałbym zrobić coś takiego:

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

Zakładam, że chcesz wykreślić różne punkty. W tym przypadku, jeżeli zdefiniować numpy tablicy:

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

następnie można wygenerować szereg kolorach 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() 

plony