Matplotlib: Как придать цвет окружностям диаграммы рассеяния на основе значений столбцов

Question

У меня есть данные, состоящие из 3 столбцов:

zone | pop1 | pop2
----   ----   ----
3      4500   3800
2      2800   3100
1      1350   1600
2      2100   1900
3      3450   3600

Я хотел бы нарисовать график рассеяния pop1 и pop2, с кружками, имеющими цвета, основанные на значении zone.

Пока у меня есть следующий код:

df = pd.read_csv(file_path)
plt.scatter(df['pop1'],df['pop2'], s = 1)

Как я могу дать разные цвета, скажем, красный, зеленый и синий, соответствующие значениям зон 1, 2 и 3 соответственно?

Answer 1

Не используя дополнительную библиотеку, вы также можете перейти на что-то вроде:

colors = {1:'red', 2:'green', 3:'blue'}

for i in range(len(df)):
    plt.scatter(df['pop1'].iloc[i], df['pop2'].iloc[i],
                c=colors[df['zone'].iloc[i]])

РЕДАКТИРОВАТЬ: вам не нужно использовать цикл, вы можете использовать что-то вроде этого:

colors = {1:'red', 2:'green', 3:'blue'}

plt.scatter(df['pop1'], df['pop2'], 
            c=[colors[i] for i in df['zone']])

Что дает вывод:

Для этого необходимо создать словарь цветов для значений в zones. Также вы потратите дополнительное время на составление списка.

Answer 2

Вы можете использовать пакет seaborn, в котором используется оболочка matplotlib. У этого есть множество особенностей с красивыми заговорами. Вот простой пример вашего вопроса.

import matplotlib.pyplot as plt
%matplotlib inline 
import seaborn as sns
import pandas as pd

data = pd.DataFrame({'col1':[4500,2800,1350,2100,3450],
             'col2':[3800,3100 ,1650,1900,3600],
             'col3':[3,2,1,2,3]})

sns.lmplot(data=data, x='col1', y='col2', hue='col3', 
                   fit_reg=False, legend=True)
#fit_reg is use to fit a line for regression, we need only dots.