pythonpandasdata-cleaningdata-preprocessing

Pandas pivot/collapse on specific conditions

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']})

enter image description here

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:

enter image description here

Solution

  • 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