Change color of matplotlib.pyplot points

Question

Here is plot code I've written :

import matplotlib.pyplot as plt

Y = [ 1 , 2 , 3 ]
X = [ 1 , 2 , 4 ]
vocabulary = [1 , 2 , 3]

plt.scatter(X , Y)
for label, x, y in zip(vocabulary, X, Y):
    if(label == 1):
        plt.annotate('', xy=(x, y), xytext=(0, 0), color='red' , textcoords='offset points')
    elif(label == 1):
        plt.annotate('', xy=(x, y), xytext=(0, 0), color='green' , textcoords='offset points')
    elif(label == 1):
        plt.annotate('', xy=(x, y), xytext=(0, 0), color='blue' , textcoords='offset points')
    else :
        plt.annotate('', xy=(x, y), xytext=(0, 0), color='black' , textcoords='offset points')
plt.show()

I'm attempting to change the color depending on the value in array vocabulary if 1 then color the data point red , if 2 then green , if 3 then blue else color the point black. But for all points the color of each point is set to blue. How to color the data point depending on the current value of vocabulary ?

Above code produces :

Answer 1

You could make a dictionary of colors and look it up during scatter plot like shown below

%matplotlib inline
import matplotlib.pyplot as plt

Y = [ 1 , 2 , 3 ,6]
X = [ 1 , 2 , 4 ,5]
vocabulary = [1 , 2 , 3, 0]
my_colors = {1:'red',2:'green',3:'blue'}

for i,j in enumerate(X):
    # look for the color based on vocabulary, if not found in vocubulary, then black is returned.
    plt.scatter(X[i] , Y[i], color = my_colors.get(vocabulary[i], 'black'))

plt.show()

results in