Легенда разброса Matplotlib с цветами с использованием категориальной переменной

Question

Я сделал простую диаграмму рассеяния, используя matplotlib, показывая данные от 2 числовых переменных (varA и varB) с цветами, которые я определил с помощью 3-й категориальной строковой переменной (col), содержащей 10 уникальных цветов (соответствующих другой строковой переменной с 10 уникальными именами ), все в том же DataFrame Pandas со 100+ строками. Есть ли простой способ создать легенду для этой диаграммы рассеяния, которая отображает уникальные цветные точки и названия соответствующих категорий? Или я должен как-то сгруппировать данные и построить каждую категорию в подзаговоре, чтобы сделать это? Это то, что я до сих пор:

import matplotlib.pyplot as plt
from matplotlib import colors as mcolors

varA = df['A']
varB = df['B'] 
col = df['Color']

plt.scatter(varA,varB, c=col, alpha=0.8)
plt.legend()

plt.show()

Answer 1

Учитывая, что Color - это столбец, который имеет все цвета и метки, вы можете просто сделать следующее.

colors = list(df['Color'].unique())
for i in range(0 , len(colors)):
    data = df.loc[df['Color'] == colors[i]]
    plt.scatter('A', 'B', data=data, color='Color', label=colors[i])
plt.legend()
plt.show()

Answer 2

Простой способ - сгруппировать данные по цвету, а затем нанести все данные на один график. Панды имеет встроенную функцию группового. Например:

import matplotlib.pyplot as plt
from matplotlib import colors as mcolors

for color, group in df.groupby(['Color']):
    plt.scatter(group['A'], group['B'], c=color, alpha=0.8, label=color)

plt.legend()
plt.show()

Обратите внимание, что мы вызываем plt.scatter один раз для каждой группировки данных. Тогда нам нужно только вызвать plt.legend и plt.show, как только все данные будут в нашем графике.