Extract values from different columns in Python-CodePudding

I am looking for a more Pythonic way to extract values from columns, that comes in pairs of two. Here is my solution so far, but I am looking for a faster (possible without iteration???) way.

>>> import pandas as pd
>>> data = {
...     "Home Team": ["A", "B", "C", "B", "A"],
...     "Away Team": ["B", "A", "A", "A", "D"],
...     "Home Feat": [1, 2, 0, 10, 7],
...     "Away Feat": [2, 3, 0, 5, 1],
... }
>>> df = pd.DataFrame(data=data)
>>> def extract(df, team):
...     result = list()
...     for index, row in df.iterrows():
...         home = row['Home Team']
...         away = row['Away Team']
...         if home == team:
...             result.append(row['Home Feat'])
...         if away == team:
...             result.append(row['Away Feat'])
...     return {team: result}
>>> extract(df, "A")
{'A': [1, 3, 0, 5, 7]}

CodePudding user response：

Solution with create MultiIndex by split columns names by space and then reshape by DataFrame.stack with filter Teams:

df.columns = df.columns.str.split(expand=True)
L = df.stack(0).loc[lambda x: x['Team'] == 'A', 'Feat'].tolist()
print (L)
[1, 3, 0, 5, 7]