matplotlib multiple line plot color according to reference value-CodePudding

If i have a dataset:

import numpy as np
import matplotlib.pyplot as plt
y = np.array([np.random.randn(10) for k in range(100)])
plt.plot(y)

And a corresponding vector of categorical values:

x = np.array([['nr1', 'nr2'] for k in range(5)]).reshape(10,1)

How do i color according to whether the y is 'nr1' or 'nr2'? It's also ok if pandas plot() function have an easy way.

It should preferably be scalable so there can be more than two groups in x.

CodePudding user response：

Edit: This already includes the comment and is working for any number of categories.

import numpy as np
import matplotlib.pyplot as plt
import matplotlib as mpl
y = np.array([np.random.randn(6) for k in range(100)])

x = np.array([['nr1', 'nr2', 'n3'] for k in range(2)]).reshape(6)

uniques, inverse = np.unique(x, return_inverse=True)
c = mpl.cm.Set1(inverse)

for i in range(y.shape[1]):
    plt.plot(y[:,i], color=c[i])