如何在 python 3 jupyter 中添加图例和颜色代码列表

Question

我正在尝试从 csv 文件创建散点图并且它有效，但我现在想做的是根据我制作的列表制作一个图例，这是我在代码中制作的国家/地区列表，我也想要图表旁边的图例。但我在网上找到的大多数解决方案都使用了 numpy 和 panda 等库，但我们不允许使用除 NamedTuple、List、matplotlib.pyplot 和 csv 之外的任何库。有办法这样做吗？以下是数据截图：

**DATA DEFINITION**

from cs103 import *
from typing import NamedTuple, List
from enum import Enum
import csv
import matplotlib.pyplot as plt

##################
# Data Definitions

#The file we are using is called TENNIS_ACE_VS_WIN.csv and it includes rank, name, country name, country ID, ace percentages, and match
#won percentages. We chose to use name, country ID, ace percentages and matches won percentages to represent our information. They are
#crucial to our project because it is all the information we need in order to achieve our research question.

TennisPlayer = NamedTuple('TennisPlayer', [('name', str),
                                           ('country_id', str),
                                           ('ace_percentage', float), #represented as a percentage, (e.g., 20.5 for 20.5%), in range [0, 100]
                                           ('matches_won_percentage', float)])  #represented as a percentage, (e.g., 20.5 for 20.5%),
                                                                                #in range [0, 100]

#interp. a tennis player including their name, country ID, aces percentage, and matches won percentage
HURKACZ = TennisPlayer("Hubert Hurkacz", "POL", 24.94, 59.44)
FRITZ = TennisPlayer("Taylor Fritz", "USA", 18.47, 58.26)
KORDA = TennisPlayer("Sebastian Korda", "USA", 12.65, 53.61)


@typecheck
def fn_for_tennis_player(tp: TennisPlayer) -> ...:
    #template based on compound
    return ...(tp.name,
               tp.country_id,
               tp.ace_percentage,
               tp.matches_won_percentage)

# List[TennisPlayer]
# interp. a list of TennisPlayer representing the statistics of professional tennis players

LOTP0 = []
LOTP1 = [HURKACZ, FRITZ]
LOTP2 = [HURKACZ, FRITZ, KORDA]

@typecheck
def fn_for_lotp(lotp: List[TennisPlayer]) -> ...:
    #description of the acc
    acc = ...  #type: ...
    for tp in lotp:
        acc = ...
        
    return ...

LOC0 = []
LOC1 = ["POL", "USA", "RUS", "KAZ"]
LOC2 = ["CAN", "USA", "CHN", "JPN"]


@typecheck
def fn_for_loc(loc: List[countries]) -> ...:
    #description of the acc
    acc = ...  #type: ...
    for c in loc:
        acc = ...
        
    return ... 



# Consumed = ...



# # List[Consumed]
# # interp. a list of Consumed

# LOC0 = []

# @typecheck
# def fn_for_loc(loc: List[Consumed]) -> ...:
#     ... # choose which template body to use for List[Consumed]


###########
# Functions

countries = ['POL', 'USA', 'RUS', 'GER', 'KAZ', 'NED', 'GRE', 'AUS', 'BUL', 'CHN', 'FRA', 'ITA', 'SUI', 'CAN', 'ARG', 'SRB', 
             'DEN', 'NOR', 'ESP', 'CZE', 'GBR', 'FIN', 'AUT', 'RSA', 'JPN', 'CRO', 'COL', 'BRA', 'CHI', 'PER','HUN', 'BLR']

@typecheck
def main(filename: str, countries: List[str]) -> None:
    """
    Reads the file from given filename, analyzes the data, returns the result 
    """
    # Template from HtDAP, based on function composition 
    return show_scatterplot(read(filename), countries)
                   

    

@typecheck
def read(filename: str) -> List[TennisPlayer]:
    """    
    reads information from 'TENNIS_ACE_VS_WIN.csv' file and returns a list of TennisPlayer data
    """
    #return []  #stub
    # Template from HtDAP
    # loc contains the result so far
    lotp = [] # type: List[TennisPlayer]

    with open(filename) as csvfile:
        
        reader = csv.reader(csvfile)
        next(reader) # skip header line

        for row in reader:
            # you may not need to store all the rows, and you may need
            # to convert some of the strings to other types
            tp = TennisPlayer(row[1],              #just a str
                              row[3],              #just a str
                              parse_float(row[4]), #represents a float, in range [0, 100]
                              parse_float(row[5])) #represents a float, in range [0, 100]

            lotp.append(tp)

    return lotp



@typecheck
def find_avg_aces(lotp: List[TennisPlayer], loc: List[str]) -> List[float]:
    """
    Takes in a list of TennisPlayer and a list of countries and returns a list of the averages of all the aces of each country
    """
    #return [] #stub
    
    #initialize the values
    acc = []
    aces = 0.0 #type: float
    count = 0 #type: int
    
    #description of the acc
    for country in loc:
        for player in lotp:
            if player.country_id == country:
                aces = aces + player.ace_percentage
                count = count + 1
        if count == 0:
            avg_aces = 0
        else:
            avg_aces = aces / count
        acc.append(avg_aces)
        
    return acc
    

@typecheck
def find_avg_matches_won(lotp: List[TennisPlayer], loc: List[str]) -> List[float]:
    """
    Takes in a list of TennisPlayer and a list of countries and returns a list of the averages of all the matches won of each country
    """
    #return [] #stub
    
    #initialize the values
    acc = []
    matches_won = 0.0 #type: float
    count = 0 #type: int
    
    #description of the acc
    for country in loc:
        for player in lotp:
            if player.country_id == country:
                matches_won = matches_won + player.matches_won_percentage
                count = count + 1
        if count == 0:
            avg_matches_won = 0
        else:
            avg_matches_won = matches_won / count
        acc.append(avg_matches_won)

    return acc


def show_scatterplot(lotp: List[TennisPlayer], loc: List[str]) -> None:
    """
    Given a list of ace percentages and a list of matches won percentages, return a scatterplot graph.
    Colour code the points by country with 32 countries. 
    """
    # return None  # stub

    x_vals = find_avg_aces(lotp,loc)
    y_vals = find_avg_matches_won(lotp,loc)
    
    # set the labels for the axes
    plt.xlabel('Ace percentage %')
    plt.ylabel('Matches won percentage %')
    plt.title('Relationship betweeen Matches Won % and Ace %')

    # create the scatterplot
    #plt.scatter(x_vals, y_vals, s = 1)

    #Colourcoding points
    scatter = plt.scatter(x_vals, y_vals, s = 1, c = 'green')

    
    # show the plot
    plt.show()
    
    return None
    

# Examples and tests for main
start_testing()
expect(main("TENNIS_ACE_VS_WIN_empty.csv", countries), None)
expect(main("TENNIS_ACE_VS_WIN_test1.csv", countries), None)
expect(main("TENNIS_ACE_VS_WIN_test2.csv", countries), None)
summary()

start_testing()
# Examples and tests for read
expect(read("TENNIS_ACE_VS_WIN_empty.csv"), [])

expect(read("TENNIS_ACE_VS_WIN_test1.csv"), [TennisPlayer("Alexei Popyrin", "AUS", 13.39, 53.95),
        TennisPlayer("Maxime Cressy", "USA", 11.23, 52.43),
        TennisPlayer("Adrian Mannarino", "FRA", 10.42, 51.82),
        TennisPlayer("Marco Cecchinato", "ITA", 5.58, 48.05)])

expect(read("TENNIS_ACE_VS_WIN_test2.csv"), [TennisPlayer("Daniil Medvedev", "RUS", 14.72, 55.01),
        TennisPlayer("John Isner", "USA", 11.81, 53.18),
        TennisPlayer("Adrian Mannarino", "FRA", 10.42, 51.82),
        TennisPlayer("Denis Shapovalov", "CAN", 6.84, 49.11)])
summary()


#examples and tests for find_avg_aces
start_testing()
expect(find_avg_aces(LOTP1, LOC1), [24.94, 15.56, 0.0, 0.0])
# expect(find_avg_aces(LOTP2, LOC2), 15.56)
summary()

#examples and tests for find_avg_matches_won
# start_testing()
# expect(find_avg_matches_won(LOTP1, FRITZ), 58.26)
# expect(find_avg_matches_won(LOTP2, FRITZ), 55.935)
# summary()


#examples and tests for show_scatterplot
# start_testing()
# expect(show_scatterplot(LOTP1), None)
# summary()
main("TENNIS_ACE_VS_WIN.csv", countries)

Answer 1

我不完全确定你在问什么，但似乎你需要一个包含图例的散点图。我对您的散点图函数进行了一些调整，以便它可以与 NamedTuples 列表一起使用。欢迎您根据需要进行其他修改。

from cs103 import *
from typing import NamedTuple, List
import matplotlib.pyplot as plt

##################
# Data Definitions

TennisPlayer = NamedTuple('TennisPlayer', [('rank', int),
                                           ('name', str),
                                           ('country_name', str),
                                           ('country_id', str),
                                           ('ace_percentage', float),
                                           ('matches_won_percentage', float)])

player1 = TennisPlayer(1, "Novak Djokovic", "Serbia", "SRB", 14.27, 81.67)
player2 = TennisPlayer(2, "Daniil Medvedev", "Russia", "RUS", 12.75, 77.88)
player3 = TennisPlayer(3, "Rafael Nadal", "Spain", "ESP", 10.94, 81.25)
player4 = TennisPlayer(4, "Alexander Zverev", "Germany", "GER", 11.61, 77.52)
player5 = TennisPlayer(5, "Stefanos Tsitsipas", "Greece", "GRE", 10.32, 74.72)
player6 = TennisPlayer(6, "Dominic Thiem", "Austria", "AUT", 10.08, 75.44)
player7 = TennisPlayer(7, "Matteo Berrettini", "Italy", "ITA", 9.54, 70.69)
player8 = TennisPlayer(8, "Diego Schwartzman", "Argentina", "ARG", 7.59, 69.16)
player9 = TennisPlayer(9, "Denis Shapovalov", "Canada", "CAN", 11.02, 67.31)

players = [player1, player2, player3, player4, player5, player6, player7, player8, player9]

def show_scatterplot(players: List[TennisPlayer]) -> None:
    ace_percentage = [player.ace_percentage for player in players]
    matches_won_percentage = [player.matches_won_percentage for player in players]
    country_names = [player.country_name for player in players]
    unique_countries = list(set(country_names))
    colors = plt.cm.get_cmap('tab10', len(unique_countries))

    for i, country in enumerate(unique_countries):
        country_players = [player for player in players if player.country_name == country]
        country_ace_percentage = [player.ace_percentage for player in country_players]
        country_matches_won_percentage = [player.matches_won_percentage for player in country_players]
        plt.scatter(country_ace_percentage, country_matches_won_percentage, color=colors(i), label=country)
    plt.xlabel("Ace Percentage")
    plt.ylabel("Matches Won Percentage")
    plt.title("Tennis Player Performance")
    plt.legend(title='Country')
    plt.show()

show_scatterplot(players)