CSV-печать в Python в соответствии с категорией / названием

Question

Мне удалось построить данные в соответствии с тем, какой номер столбца кто-то вводит для оси x и y.Мои данные отформатированы как CSV, где в пятой колонке есть название вида:

5.1,3.5,1.4,0.2, Iris-setosa

7.0,3.2,4.7,1.4, Iris-versicolor

5.8,2.7,5.1,1.9, Iris-virginica

В данный момент моя программа работает правильно и отображает точки, в которых они мне нужны.Проблема в том, что все точки одного цвета.Мне нужно как-то сказать программе, чтобы она смотрела на названия видов и использовала их в качестве категорий для соответствующих значений.Существует три вида данных, поэтому данные должны быть в трех цветах с легендой.

import random
import matplotlib.pyplot as plt
import csv

#a function that will take data from a CSV and plot them according to which columns are inputted

def plot_data(fileName,colX,colY):
    dataList = []
    sepalLengthCM = []
    sepalWidthCM= []
    petalLengthCM =[]
    petalWidthCM = []
    species = []

    #reading the file
    with open(fileName, "r") as file:
        data = csv.reader(file) 

        #making a list of all the rows of data
        for row in data:
            dataList.append(row)

        #seperating each column into it's own list so I can plot them against eachother. For example, I'm plotting row 2 as the x axis and row 1 as the y
    for row in range(0, len(dataList)-1):
        sepalLengthCM.append(dataList[row][0])
        sepalWidthCM.append(dataList[row][1])
        petalLengthCM.append(dataList[row][2])
        petalWidthCM.append(dataList[row][3])
        species.append(dataList[row][4])

    #placing each column into a list of 'options' that the user can choose from.
    optionsList = [sepalLengthCM, sepalWidthCM, petalLengthCM, petalWidthCM]
    #using the indexes of the options list to plot the scatter plot. It works, but without distinction among species
    plt.scatter(optionsList[colX],optionsList[colY])
    plt.show()


plot_data("iris.csv",2,1)

Как мне сказать Python посмотреть на этот четвертый столбец?Я разделил название вида на его собственный список, но не думаю, что оно мне здесь нужно.Я знаю, как строить столбцы, но не могу понять, как классифицировать строки.

Answer 1

Вам просто нужно использовать аргумент c функции scatter, здесь приведен пример

SPECIES_COLORS = {
    "Iris-setosa": "b",
    "Iris-versicolor": "y",
    "Iris-virginica": "r",
}

colors = [SPECIES_COLORS[s] for s in species]

plt.scatter(optionsList[colX],optionsList[colY],c=colors)

Конечно, словарь SPECIES_COLORS можно определить вне тела функции.

источник: документация по функции рассеяния