Pandas pivot/collapse on specific conditions

Question

I have this dataset:

df = pd.DataFrame({'game_id' : [123,123,456,456],
                   'location' : ['home', 'away','home', 'away'],
                   'away_team' : ['braves', 'braves', 'mets', 'mets'],
                   'home_team' : ['phillies', 'phillies', 'marlins', 'marlins']})

I want to transform it so that the away_team and home_team columns are collapsed into 1 "team" column that aligns with location, for each "game_id". I would want it to look like this:

Answer 1

Try:

df["team"] = df.apply(lambda x: x[x["location"] + "_team"], axis=1)
print(df[["game_id", "location", "team"]])

Prints:

   game_id location      team
0      123     home  phillies
1      123     away    braves
2      456     home   marlins
3      456     away      mets

Answer 2

Use melt to pivot your data and then use boolean indexing filter out the rows where the location does not align with the home_team or away_team column (i.e., variable)

# melt you frame
df_m = df.melt(['game_id', 'location'], ['away_team', 'home_team'])
# boolean indexing to keep only the values that match the location
# you can then drop the variable column and sort if you want
df_m[df_m['location'].eq(df_m['variable'].str[:4])].drop(columns='variable').sort_values('game_id')

   game_id location     value
1      123     away    braves
4      123     home  phillies
3      456     away      mets
6      456     home   marlins

Answer 3

A one liner using melt:

out = (df.rename(columns={'away_team': 'away', 'home_team': 'home'})
         .melt(['game_id', 'location'], value_name='team')
         .loc[lambda x: x['location'] == x.pop('variable')]
         .sort_values(['game_id', 'location'], ascending=[True, False], ignore_index=True))

Output:

>>> out
   game_id location      team
0      123     home  phillies
1      123     away    braves
2      456     home   marlins
3      456     away      mets