Как основать сегмент цветовой линии на значении строки / другого столбца?

Question

фрейм данных с 3 столбцами

x       y     type
21/02   5     'a'
22/02   6     'b'
...   ...   ...

всего у меня есть два типа: 'a' and 'b'

на основе stackoverflow вопрос и документация matplotlib я пришел со следующим:

fig, ax = plt.subplots()   
xy = np.column_stack((df['x'],df['y'])) 
xy = xy.reshape(-1, 1, 2) 
segments = np.hstack([xy[:-1], xy[1:]]) 
coll = LineCollection(segments, color='r') 
ax.add_collection(coll)
plt.show()

Это дает мне следующий график:

Я думаю, что должен что-то сделатьс: coll.set_array(some_value) Но я не знаю как.

Я получил это далеко: я использовал тот же код, что и выше, но только изменил одну строку:

coll = LineCollection(segments,cmap=plt.cm.gist_ncar)

Но теперьЦвет сегментов наугад.

Как мне основать цвет сегментов на столбце type в моем фрейме данных?

Answer 1

Если у вас есть только два цвета, это должно работать:

colors = ["red" if type=="a" else "blue" for type in df["type"]] 
coll = LineCollection(segments, color=colors)

Для большего количества цветов вы можете сопоставить типы со списком цветов аналогичным образом.