Matplotlib. Изменить цвет линии на основе значений в другом столбце, сохранив ту же ось х

Question

У меня следующий набор данных:

import pandas as pd
import matplotlib.pyplot as plt
dict = {'time':["2017-01-02", "2017-01-03", "2017-01-04", "2017-01-05", "2017-01-06"],'val':[3.2, 10.2, 11.3, 4.9, 2.3], 
        'class': [0, 1, 1, 0,0]} 
df = pd.DataFrame(dict) 
df
    time         val    class
0   2017-01-02   3.2    0
1   2017-01-03   10.2   1
2   2017-01-04   11.3   1
3   2017-01-05   4.9    0
4   2017-01-06   2.3    0

Я хочу построить линию для столбца «val», сохранив ось x как «df.time», тем временем изменив цвет линии на основе столбца «class» ( когда он равен нулю, то, например, синяя линия, когда он равен 1, он меняет цвет на красный). мой график выглядит следующим образом:

, но желательно что-то вроде этого:

Спасибо!

Answer 1

Как и в этот вопрос , вам просто нужно построить несколько строк:

# recommend
df['time'] = pd.to_datetime(df['time'])

plt.figure(figsize=(10,6))
for i in range(1,len(df)):
    s = df.iloc[i-1:i+1]
    color = 'r' if s['class'].eq(1).all() else 'C0'

    plt.plot(s['time'], s['val'], c=color)

plt.show()

Вывод:

Поскольку, когда у вас много строк, может быть лучше использовать scatter:

plt.scatter(df['time'], df['val'], 
            color=np.where(df['class'], 'r','C0')
           )

Вывод (будет выглядеть лучше с 10k строками):