Добавление столбца в df с функцией

Question

      Date             Visitor  V_PTS                 Home  H_PTS  \
0 2012-10-30 19:00:00  Washington Wizards     84  Cleveland Cavaliers     94   
1 2012-10-30 19:30:00    Dallas Mavericks     99   Los Angeles Lakers     91   
2 2012-10-30 20:00:00      Boston Celtics    107           Miami Heat    120   
3 2012-10-31 19:00:00    Sacramento Kings     87        Chicago Bulls     93   
4 2012-10-31 19:30:00     Houston Rockets    105      Detroit Pistons     96   

Попытка добавить в очищенный набор данных для анализа посещаемости игр NBA. Я пытаюсь добавить несколько колонок, таких как игра на арене, а также вместимость. Вот часть функции, которую я написал, чтобы добавить арену. Есть лучший способ сделать это? У меня есть даты в datetime, так как мне правильно выбрать год, чтобы назначить правильную арену командам, которые построили новые арены за последние несколько лет (Sacramento Kings). Также есть ли способ добавить стадион к этому и убить двух зайцев одним выстрелом вместо создания другой функции?

def label_arena (hometeam):
    if hometeam == 'Toronto Raptors' :
        return 'Air Canada Centre'
    if hometeam == 'Miami Heat' :
        return 'American Airlines Arena'
    if hometeam == 'Dallas Mavericks' :
        return 'American Airlines Center'
    if hometeam == 'Orlando Magic' :
        return 'Amway Center'
    if hometeam == 'San Antonio Spurs' :
        return 'AT&T Center'
    if hometeam == 'Indiana Pacers' :
        return 'Bankers Life Fieldhouse'
    if hometeam == 'Brooklyn Nets' :
        return 'Barclays Center'
    if hometeam == 'Milwaukee Bucks' :
        return 'Bradley Center'
    if hometeam == 'Washington Wizards' :
        return 'Capital One Arena'
    if hometeam == 'Oklahoma City Thunder' :
        return 'Chesapeake Energy Arena'
    if hometeam == 'Memphis Grizzlies' :
        return 'FedExForum'
    if hometeam == 'Sacramento Kings' and df['Date'] < 2016:
        return 'Sleep Train Arena'
    if hometeam == 'Sacramento Kings' and df['Date'] > 2016:
        return 'Golden 1 Center'

Answer 1

import pandas as pd

home_arenas_capacities = pd.DataFrame([
     ['Toronto Raptors', 'Air Canada Centre', 20511],
     ['Miami Heat', 'American Airlines Arena', 19600],
     ...
    ]) 

df.merge(home_arenas_capacities, on='Home')

Для королей Сакраменто вы хотите объединить «Дом» и «Дата»> 2016, что может потребовать, чтобы вы создали временный столбец, затем df.merge(..., on=['Home','Date_GE_2016']) и удалите столбец «Дата_Г_2016».

Но более чистым способом было бы добавить столбец "Сезон" = "2015-16", "2016-17". По мере того, как ваша база данных увеличивается, кажется, вам это нужно. (Для базы данных игры вы можете автоматически извлечь «Сезон» из значения «Дата». Для фрейма данных «home_arenas_capacities» вам нужно будет отредактировать его вручную).

Answer 2

Вот метод, использующий numpy.select, если вы не против numpy:

import numpy as np

conditions = [
    df['Home'] == 'Toronto Raptors',
    df['Home'] == 'Miami Heat',
    df['Home'] == 'Dallas Mavericks',
    ...
    (df['Home'] == 'Sacramento Kings') & (df['Date'].dt.year < 2016),
    (df['Home'] == 'Sacramento Kings') & (df['Date'].dt.year > 2016)]

choices = [
        'Air Canada Centre',
        'American Airlines Arena',
        'American Airlines Center',
        ...
        'Sleep Train Arena',
        'Golden 1 Center']

df['arena'] = np.select(conditions, choices)

Обратите внимание, что для того, чтобы ваше условие df['Date'] работало, вам нужно установить df['Date'] в серию даты и времени (вы можете сделать это через df['Date'] = pd.to_datetime(df['Date']), если вы еще этого не сделали)

Answer 3

Вот что вы можете сделать, чтобы упростить свою логику:

import pandas as pd

df = pd.DataFrame({'Date': ['2012-10-30', '2012-10-30', '2012-10-30',
                            '2012-10-31', '2017-10-31'],
                   'Home': ['Toronto Raptors', 'Los Angeles Lakers', 'Miami Heat',
                            'Sacramento Kings', 'Sacramento Kings']})

df['Date'] = pd.to_datetime(df['Date'])

d = {'Toronto Raptors': 'Air Canada Centre',
     'Los Angeles Lakers': 'Staples Center',
     'Miami Heat': 'American Airlines Arena'}

# general criteria
df['Arena'] = df['Home'].map(d)

# custom criteria
df.loc[(df['Home'] == 'Sacramento Kings') &
       (df['Date'].dt.year < 2016), 'Arena'] = 'Sleep Train Arena'
df.loc[(df['Home'] == 'Sacramento Kings') &
       (df['Date'].dt.year >= 2016), 'Arena'] = 'Golden 1 Center'

print(df)

        Date                Home                    Arena
0 2012-10-30     Toronto Raptors        Air Canada Centre
1 2012-10-30  Los Angeles Lakers           Staples Center
2 2012-10-30          Miami Heat  American Airlines Arena
3 2012-10-31    Sacramento Kings        Sleep Train Arena
4 2017-10-31    Sacramento Kings          Golden 1 Center