Search code examples
pythonnumpyplotuniquelegend

Trouble with a plot legend

I have two arrays, x and y, for plotting, and a third array, z, that identifies the x and y points. The integers in z are repeated, so I made a z2 array that identifies the unique values. I need to make a plot that shows a legend from the z2 array, with the plotted points reflecting those same colors. But instead I get all one color in the plot and different colors in the legend. Here is my code.

import matplotlib.pyplot as plt
import numpy as np

x = [0.54638897, 0.74436089, 0.36840323, 0.67932601, 0.56410781, 0.20797502,
 0.54681392, 0.47598874, 0.33771962, 0.6626352,  0.06115377, 0.37277143,
 0.43410935, 0.97386762, 0.69819935, 0.62578862, 0.15594451, 0.43509243,
 0.3712351,  0.94039755]
y = [0.45281763, 0.85509999, 0.65361185, 0.87928696, 0.00333544, 0.92478824,
 0.95129375, 0.15493552, 0.06571068, 0.31728336, 0.58555545, 0.52413135,
 0.43512262, 0.91267715, 0.56997665, 0.93413675, 0.57615435, 0.18518019,
 0.98207871, 0.99850326]
z = [1,1,1,1,5,5,5,11,11,11,1,1,6,6,8,8,11,9,9]
z2 = np.unique(z)

print(z2)

for i in (z2):
    plt.plot(x, y, 'o', label=i)
    
plt.plot(x, y, 'o')
plt.legend()
plt.grid()

And this is the plot I get.

enter image description here

I need, for example, x and y values [0 through 3] to correspond to z = 1 in the plot. According to the legend, each of those dots would be colored blue. I know I'm doing something wrong here. Any advice would be appreciated.

Solution

  • You don't need to use np.unique at all. Instead, use plt.scatter and use the c argument, like so:

    plt.scatter(x, y, c=z, marker='o')

    It also appears that z has 19 elements, while x and y have 20 each.